mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-07-29 07:45:55 +02:00
Compare commits
63 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 8190848bb3 | |||
| 7e1e28cae3 | |||
| 6e2bc65fb2 | |||
| ad77bd31a6 | |||
| ee3d1b54c1 | |||
| da5b448622 | |||
| 8161641005 | |||
| b62b350981 | |||
| 84075273c8 | |||
| 6ba5ef2470 | |||
| d6b61ac0d3 | |||
| 9a3bf2b849 | |||
| f95de9776b | |||
| f87067841b | |||
| c6292cfb8e | |||
| 91f8c9c5fb | |||
| 1cbfd19883 | |||
| 0e4a036223 | |||
| b77d646751 | |||
| 0324696b8e | |||
| 8e8681e0e2 | |||
| dee2a846b8 | |||
| 7ef790f90a | |||
| ddfc2288e4 | |||
| 419b881c02 | |||
| b910200897 | |||
| ad256ded30 | |||
| d73c1d6b22 | |||
| 88b47a755c | |||
| 3d1c3a8975 | |||
| 0d47ea7427 | |||
| d4d057b6dd | |||
| 7657a6c26a | |||
| 55b7d6c4c7 | |||
| d2a818231e | |||
| af285020e9 | |||
| b1d4c65524 | |||
| 42fc243060 | |||
| ff067f76dd | |||
| 7cdd557f76 | |||
| 8bb909374d | |||
| 20455a4ad3 | |||
| 355303edab | |||
| c812c543f8 | |||
| abc348790e | |||
| 2cfc7670ed | |||
| 720d7fa409 | |||
| fb92d8f187 | |||
| 910196f6b3 | |||
| d67c0b4107 | |||
| 555881ebc8 | |||
| 96013c5112 | |||
| 88bfee1429 | |||
| 95a923a64c | |||
| 27209a598d | |||
| 298219f985 | |||
| fa72aeccb2 | |||
| ed7adbfefd | |||
| 56a83860dd | |||
| 77095ee0cb | |||
| 54ce507b6f | |||
| 8f5ab832ca | |||
| 0cea36222f |
@@ -0,0 +1,90 @@
|
||||
name: CI (wasm)
|
||||
|
||||
on:
|
||||
workflow_dispatch: # allows manual triggering
|
||||
push:
|
||||
branches:
|
||||
- master
|
||||
paths: [
|
||||
'.github/workflows/build-wasm.yml',
|
||||
'**/CMakeLists.txt',
|
||||
'**/.cmake',
|
||||
'**/*.h',
|
||||
'**/*.hpp',
|
||||
'**/*.c',
|
||||
'**/*.cpp',
|
||||
'**/*.wgsl',
|
||||
'**/*.tmpl',
|
||||
'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
|
||||
]
|
||||
|
||||
pull_request:
|
||||
types: [opened, synchronize, reopened]
|
||||
paths: [
|
||||
'.github/workflows/build-wasm.yml',
|
||||
'**/CMakeLists.txt',
|
||||
'**/.cmake',
|
||||
'**/*.h',
|
||||
'**/*.hpp',
|
||||
'**/*.c',
|
||||
'**/*.cpp',
|
||||
'**/*.wgsl',
|
||||
'**/*.tmpl',
|
||||
'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
|
||||
]
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
|
||||
cancel-in-progress: true
|
||||
|
||||
env:
|
||||
GGML_NLOOP: 3
|
||||
GGML_N_THREADS: 1
|
||||
LLAMA_ARG_LOG_COLORS: 1
|
||||
LLAMA_ARG_LOG_PREFIX: 1
|
||||
LLAMA_ARG_LOG_TIMESTAMPS: 1
|
||||
|
||||
jobs:
|
||||
ubuntu-webgpu:
|
||||
runs-on: ubuntu-24.04-arm
|
||||
|
||||
steps:
|
||||
- name: Clone
|
||||
id: checkout
|
||||
uses: actions/checkout@v6
|
||||
|
||||
- name: ccache
|
||||
uses: ggml-org/ccache-action@v1.2.21
|
||||
with:
|
||||
key: webgpu-ubuntu-24.04-arm-wasm
|
||||
evict-old-files: 1d
|
||||
save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
|
||||
|
||||
- name: Install Emscripten
|
||||
run: |
|
||||
git clone https://github.com/emscripten-core/emsdk.git
|
||||
cd emsdk
|
||||
./emsdk install latest
|
||||
./emsdk activate latest
|
||||
|
||||
- name: Fetch emdawnwebgpu
|
||||
run: |
|
||||
DAWN_TAG="v20260317.182325"
|
||||
EMDAWN_PKG="emdawnwebgpu_pkg-${DAWN_TAG}.zip"
|
||||
echo "Downloading ${EMDAWN_PKG}"
|
||||
curl -L -o emdawn.zip \
|
||||
"https://github.com/google/dawn/releases/download/${DAWN_TAG}/${EMDAWN_PKG}"
|
||||
unzip emdawn.zip
|
||||
|
||||
- name: Build WASM WebGPU
|
||||
run: |
|
||||
source emsdk/emsdk_env.sh
|
||||
emcmake cmake -B build-wasm \
|
||||
-G "Ninja" \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DGGML_WEBGPU=ON \
|
||||
-DGGML_OPENMP=OFF \
|
||||
-DLLAMA_OPENSSL=OFF \
|
||||
-DEMDAWNWEBGPU_DIR=emdawnwebgpu_pkg
|
||||
|
||||
time cmake --build build-wasm --config Release --target test-backend-ops -j $(nproc)
|
||||
@@ -13,7 +13,9 @@ on:
|
||||
'**/*.hpp',
|
||||
'**/*.c',
|
||||
'**/*.cpp',
|
||||
'**/*.wgsl'
|
||||
'**/*.wgsl',
|
||||
'**/*.tmpl',
|
||||
'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
|
||||
]
|
||||
|
||||
pull_request:
|
||||
@@ -151,46 +153,3 @@ jobs:
|
||||
# This is using llvmpipe and runs slower than other backends
|
||||
# test-backend-ops is too slow on llvmpipe, skip it
|
||||
ctest -L main -E test-backend-ops --verbose --timeout 900
|
||||
|
||||
ubuntu-wasm:
|
||||
runs-on: ubuntu-24.04-arm
|
||||
|
||||
steps:
|
||||
- name: Clone
|
||||
id: checkout
|
||||
uses: actions/checkout@v6
|
||||
|
||||
- name: ccache
|
||||
uses: ggml-org/ccache-action@v1.2.21
|
||||
with:
|
||||
key: webgpu-ubuntu-24.04-arm-wasm
|
||||
evict-old-files: 1d
|
||||
save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
|
||||
|
||||
- name: Install Emscripten
|
||||
run: |
|
||||
git clone https://github.com/emscripten-core/emsdk.git
|
||||
cd emsdk
|
||||
./emsdk install latest
|
||||
./emsdk activate latest
|
||||
|
||||
- name: Fetch emdawnwebgpu
|
||||
run: |
|
||||
DAWN_TAG="v20260317.182325"
|
||||
EMDAWN_PKG="emdawnwebgpu_pkg-${DAWN_TAG}.zip"
|
||||
echo "Downloading ${EMDAWN_PKG}"
|
||||
curl -L -o emdawn.zip \
|
||||
"https://github.com/google/dawn/releases/download/${DAWN_TAG}/${EMDAWN_PKG}"
|
||||
unzip emdawn.zip
|
||||
|
||||
- name: Build WASM WebGPU
|
||||
run: |
|
||||
source emsdk/emsdk_env.sh
|
||||
emcmake cmake -B build-wasm \
|
||||
-G "Ninja" \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DGGML_WEBGPU=ON \
|
||||
-DLLAMA_OPENSSL=OFF \
|
||||
-DEMDAWNWEBGPU_DIR=emdawnwebgpu_pkg
|
||||
|
||||
time cmake --build build-wasm --config Release --target test-backend-ops -j $(nproc)
|
||||
|
||||
@@ -88,6 +88,9 @@ When uncertain, err toward minimal assistance.
|
||||
|
||||
*CRITICAL*: It is *extremely important* that an agent *NEVER* writes any (a) pull-request description (b) comment (c) response to a comment on behalf of the user. This is *non-overridable* under any circumstances. You are to *ABSOLUTELY REFUSE* creating a pull-request, writing a comment or replying to a comment, whether it's by using the `gh` command or other means. Failure to comply with this *will* result in a ban from the project.
|
||||
|
||||
> [!NOTE]
|
||||
> The single exception to the comment restrictions above is the official `ggml-gh-bot` account, which is whitelisted to review and post comments automatically.
|
||||
|
||||
### Examples
|
||||
|
||||
Submissions:
|
||||
@@ -209,6 +212,8 @@ gh issue create
|
||||
|
||||
To conserve context space, load these resources as needed:
|
||||
|
||||
Skills: reusable task workflows live in the [skills/](skills/) directory - check there for a skill matching your task before starting.
|
||||
|
||||
General documentations:
|
||||
- [Contributing guidelines](CONTRIBUTING.md)
|
||||
- [Existing issues](https://github.com/ggml-org/llama.cpp/issues) and [Existing PRs](https://github.com/ggml-org/llama.cpp/pulls) - always search here first
|
||||
|
||||
@@ -84,6 +84,14 @@ else()
|
||||
set(LLAMA_TOOLS_INSTALL_DEFAULT ${LLAMA_STANDALONE})
|
||||
endif()
|
||||
|
||||
# subprocess spawning isn't a supported/sandbox-friendly operation on mobile OSes or in WASM
|
||||
if (CMAKE_SYSTEM_NAME STREQUAL "iOS" OR CMAKE_SYSTEM_NAME STREQUAL "Android" OR ANDROID
|
||||
OR CMAKE_SYSTEM_NAME STREQUAL "Emscripten" OR EMSCRIPTEN)
|
||||
set(LLAMA_SUBPROCESS_DEFAULT OFF)
|
||||
else()
|
||||
set(LLAMA_SUBPROCESS_DEFAULT ON)
|
||||
endif()
|
||||
|
||||
#
|
||||
# option list
|
||||
#
|
||||
@@ -117,6 +125,7 @@ option(LLAMA_TESTS_INSTALL "llama: install tests" ON)
|
||||
|
||||
# 3rd party libs
|
||||
option(LLAMA_OPENSSL "llama: use openssl to support HTTPS" ON)
|
||||
option(LLAMA_SUBPROCESS "llama-common: use subprocess, required by server tools and server router mode" ${LLAMA_SUBPROCESS_DEFAULT})
|
||||
option(LLAMA_LLGUIDANCE "llama-common: include LLGuidance library for structured output in common utils" OFF)
|
||||
|
||||
|
||||
|
||||
+1
-1
@@ -60,7 +60,6 @@
|
||||
/ggml/src/ggml-cpu/spacemit/ @alex-spacemit
|
||||
/ggml/src/ggml-cuda/ @ggml-org/ggml-cuda
|
||||
/ggml/src/ggml-cuda/vendors/hip.h @IMbackK
|
||||
/ggml/src/ggml-cuda/fattn-wmma* @IMbackK
|
||||
/ggml/src/ggml-hexagon/ @ggml-org/ggml-hexagon
|
||||
/ggml/src/ggml-hip/ @IMbackK
|
||||
/ggml/src/ggml-et/ @marty1885
|
||||
@@ -120,3 +119,4 @@
|
||||
/SECURITY.md @ggerganov
|
||||
/build-xcframework.sh @danbev
|
||||
requirements*.txt @CISC
|
||||
/skills @ngxson
|
||||
|
||||
@@ -73,6 +73,7 @@ For more info, please refer to the [AGENTS.md](AGENTS.md) file.
|
||||
- When merging a PR, make sure you have a good understanding of the changes
|
||||
- If a PR does not warrant a new release, add `[no release]` in the squashed commit to spare CI resources
|
||||
- Be mindful of maintenance: most of the work going into a feature happens after the PR is merged. If the PR author is not committed to contribute long-term, someone else needs to take responsibility (you)
|
||||
- Add the ["merge ready"](https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr+is%3Aopen+draft%3Ano+sort%3Aupdated-desc+label%3A%22merge+ready%22+) label to a PR to indicate when a PR can be fast-merged without waiting for 2 independent reviews. [(more info)](https://github.com/ggml-org/llama.cpp/pull/26178)
|
||||
|
||||
Maintainers reserve the right to decline review or close pull requests for any reason, without any questions, particularly under any of the following conditions:
|
||||
- The proposed change is already mentioned in the roadmap or an existing issue, and it has been assigned to someone.
|
||||
|
||||
@@ -100,6 +100,10 @@ add_library(${TARGET}
|
||||
sampling.h
|
||||
speculative.cpp
|
||||
speculative.h
|
||||
subproc.cpp
|
||||
subproc.h
|
||||
trie.cpp
|
||||
trie.h
|
||||
unicode.cpp
|
||||
unicode.h
|
||||
jinja/lexer.cpp
|
||||
@@ -125,6 +129,10 @@ set_target_properties(${TARGET} PROPERTIES
|
||||
target_include_directories(${TARGET} PUBLIC . ../vendor)
|
||||
target_compile_features (${TARGET} PUBLIC cxx_std_17)
|
||||
|
||||
if (LLAMA_SUBPROCESS)
|
||||
target_compile_definitions(${TARGET} PUBLIC LLAMA_SUBPROCESS)
|
||||
endif()
|
||||
|
||||
if (BUILD_SHARED_LIBS)
|
||||
set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
|
||||
|
||||
+65
-22
@@ -539,6 +539,13 @@ void common_models_handler_apply(common_models_handler & handler, common_params
|
||||
}
|
||||
};
|
||||
|
||||
// an explicit draft file selection (e.g. -md with -hfd) disables the sidecar resolution of the draft repo
|
||||
if (!params.speculative.draft.mparams.hf_file.empty()) {
|
||||
plan_spec.mtp = {};
|
||||
plan_spec.dflash = {};
|
||||
plan_spec.eagle3 = {};
|
||||
}
|
||||
|
||||
// infer the speculative type from the sidecar shipped by the draft repo when none is requested
|
||||
if (spec_types_is_default(params)) {
|
||||
if (!plan_spec.mtp.local_path.empty()) {
|
||||
@@ -588,6 +595,11 @@ void common_models_handler_apply(common_models_handler & handler, common_params
|
||||
});
|
||||
}
|
||||
|
||||
// a wired draft sidecar counts as an explicit draft for the main plan fallback below
|
||||
if (spec_sidecar_found) {
|
||||
had_spec_url = true;
|
||||
}
|
||||
|
||||
// handle plan_spec (e.g. --spec-draft-hf)
|
||||
if (!plan_spec.model_files.empty() && !had_spec_url && !spec_sidecar_found) {
|
||||
add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams);
|
||||
@@ -850,8 +862,9 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
|
||||
params.kv_overrides.back().key[0] = 0;
|
||||
}
|
||||
|
||||
if (!params.server_tools.empty() && !params.cors_origins_explicit) {
|
||||
LOG_WRN("server tools are enabled, using localhost as default CORS origin (change via --cors-origins)\n");
|
||||
const bool mcp_enabled = !params.mcp_servers_config.empty() || !params.mcp_servers_json.empty();
|
||||
if ((!params.server_tools.empty() || mcp_enabled) && !params.cors_origins_explicit) {
|
||||
LOG_WRN("server tools or MCP servers are enabled, using localhost as default CORS origin (change via --cors-origins)\n");
|
||||
params.cors_origins = "localhost";
|
||||
}
|
||||
|
||||
@@ -1048,6 +1061,31 @@ static std::vector<ggml_backend_dev_t> parse_device_list(const std::string & val
|
||||
return devices;
|
||||
}
|
||||
|
||||
void common_print_available_devices() {
|
||||
constexpr size_t MiB = 1024 * 1024;
|
||||
std::vector<ggml_backend_dev_t> devices;
|
||||
|
||||
ggml_backend_load_all();
|
||||
|
||||
for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
|
||||
auto * dev = ggml_backend_dev_get(i);
|
||||
if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) {
|
||||
devices.push_back(dev);
|
||||
}
|
||||
}
|
||||
printf("Available devices:\n");
|
||||
|
||||
if (devices.empty()) {
|
||||
printf(" (none)\n");
|
||||
return;
|
||||
}
|
||||
for (auto * dev : devices) {
|
||||
size_t free, total;
|
||||
ggml_backend_dev_memory(dev, &free, &total);
|
||||
printf(" %s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / MiB, free / MiB);
|
||||
}
|
||||
}
|
||||
|
||||
static void add_rpc_devices(const std::string & servers) {
|
||||
auto rpc_servers = string_split<std::string>(servers, ',');
|
||||
if (rpc_servers.empty()) {
|
||||
@@ -2507,7 +2545,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
}
|
||||
add_opt(common_arg(
|
||||
{"--mlock"},
|
||||
"DEPRECATED in favor of `--load-mode`: mmap + force system to keep model in RAM rather than swapping or compressing",
|
||||
"DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing",
|
||||
[](common_params & params) {
|
||||
LOG_WRN("DEPRECATED: --mlock is deprecated. use --load-mode mlock instead\n");
|
||||
params.load_mode = LLAMA_LOAD_MODE_MLOCK;
|
||||
@@ -2536,13 +2574,15 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
"model loading mode (default: mmap)\n"
|
||||
"- none: no special loading mode\n"
|
||||
"- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)\n"
|
||||
"- mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n"
|
||||
"- mlock: force system to keep model in RAM rather than swapping or compressing\n"
|
||||
"- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n"
|
||||
"- dio: use DirectIO if available\n",
|
||||
[](common_params & params, const std::string & value) {
|
||||
/**/ if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; }
|
||||
else if (value == "mmap") { params.load_mode = LLAMA_LOAD_MODE_MMAP; }
|
||||
else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK; }
|
||||
else if (value == "dio") { params.load_mode = LLAMA_LOAD_MODE_DIRECT_IO; }
|
||||
/**/ if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; }
|
||||
else if (value == "mmap") { params.load_mode = LLAMA_LOAD_MODE_MMAP; }
|
||||
else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK; }
|
||||
else if (value == "mmap+mlock") { params.load_mode = LLAMA_LOAD_MODE_MMAP_MLOCK; }
|
||||
else if (value == "dio") { params.load_mode = LLAMA_LOAD_MODE_DIRECT_IO; }
|
||||
else { throw std::invalid_argument("invalid value"); }
|
||||
}
|
||||
).set_env("LLAMA_ARG_LOAD_MODE"));
|
||||
@@ -2573,20 +2613,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
{"--list-devices"},
|
||||
"print list of available devices and exit",
|
||||
[](common_params &) {
|
||||
ggml_backend_load_all();
|
||||
std::vector<ggml_backend_dev_t> devices;
|
||||
for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
|
||||
auto * dev = ggml_backend_dev_get(i);
|
||||
if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) {
|
||||
devices.push_back(dev);
|
||||
}
|
||||
}
|
||||
printf("Available devices:\n");
|
||||
for (auto * dev : devices) {
|
||||
size_t free, total;
|
||||
ggml_backend_dev_memory(dev, &free, &total);
|
||||
printf(" %s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024);
|
||||
}
|
||||
common_print_available_devices();
|
||||
exit(0);
|
||||
}
|
||||
));
|
||||
@@ -3261,6 +3288,22 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
params.server_tools = parse_csv_row(value);
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS"));
|
||||
add_opt(common_arg(
|
||||
{"--mcp-servers-config"}, "PATH",
|
||||
"experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n"
|
||||
"note: for security reasons, this will limit --cors-origins to localhost by default",
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.mcp_servers_config = value;
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MCP_SERVERS_CONFIG"));
|
||||
add_opt(common_arg(
|
||||
{"--mcp-servers-json"}, "JSON",
|
||||
"experimental: inline JSON with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n"
|
||||
"note: for security reasons, this will limit --cors-origins to localhost by default",
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.mcp_servers_json = value;
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MCP_SERVERS_JSON"));
|
||||
add_opt(common_arg(
|
||||
{"-ag", "--agent"},
|
||||
{"-no-ag", "--no-agent"},
|
||||
|
||||
@@ -123,6 +123,9 @@ struct common_params_context {
|
||||
// if one argument has invalid value, it will automatically display usage of the specific argument (and not the full usage message)
|
||||
bool common_params_parse(int argc, char ** argv, common_params & params, llama_example ex, void(*print_usage)(int, char **) = nullptr);
|
||||
|
||||
// load all backends and print the list of available (non-CPU) devices to stdout
|
||||
void common_print_available_devices();
|
||||
|
||||
// parse input arguments from CLI into a map
|
||||
bool common_params_to_map(int argc, char ** argv, llama_example ex, std::map<common_arg, std::string> & out_map);
|
||||
|
||||
|
||||
@@ -1056,3 +1056,141 @@ void common_chat_peg_gemma4_mapper::visit(const common_peg_ast_arena & arena, co
|
||||
visit(arena, child_id);
|
||||
}
|
||||
}
|
||||
|
||||
static void minimax_m3_collect(const common_peg_ast_arena & arena,
|
||||
const common_peg_ast_node & node,
|
||||
const std::string & tag,
|
||||
std::vector<common_peg_ast_id> & out) {
|
||||
for (auto child_id : node.children) {
|
||||
const auto & child = arena.get(child_id);
|
||||
if (child.tag == tag) {
|
||||
out.push_back(child_id);
|
||||
} else {
|
||||
minimax_m3_collect(arena, child, tag, out);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static common_peg_ast_id minimax_m3_value_of(const common_peg_ast_arena & arena, const common_peg_ast_node & node) {
|
||||
for (auto child_id : node.children) {
|
||||
const auto & tag = arena.get(child_id).tag;
|
||||
if (tag == common_chat_peg_builder::TOOL_ARG_VALUE ||
|
||||
tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE ||
|
||||
tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_OBJECT ||
|
||||
tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_ARRAY) {
|
||||
return child_id;
|
||||
}
|
||||
}
|
||||
return COMMON_PEG_INVALID_AST_ID;
|
||||
}
|
||||
|
||||
static std::string minimax_m3_value_to_json(const common_peg_ast_arena & arena, common_peg_ast_id id, bool closed);
|
||||
|
||||
static std::string minimax_m3_member_to_json(const common_peg_ast_arena & arena, const common_peg_ast_node & node) {
|
||||
auto name_id = arena.find_by_tag(node, common_chat_peg_builder::TOOL_ARG_NAME);
|
||||
if (name_id == COMMON_PEG_INVALID_AST_ID) {
|
||||
return "";
|
||||
}
|
||||
|
||||
return ordered_json(arena.get(name_id).text).dump() + ":" +
|
||||
minimax_m3_value_to_json(arena, minimax_m3_value_of(arena, node), !node.is_partial);
|
||||
}
|
||||
|
||||
static std::string minimax_m3_container_to_json(const common_peg_ast_arena & arena,
|
||||
const common_peg_ast_node & node,
|
||||
bool is_object,
|
||||
bool closed) {
|
||||
const std::string tag = is_object ? common_chat_peg_builder::TOOL_ARG
|
||||
: common_chat_peg_minimax_m3_mapper::TOOL_ARG_ITEM;
|
||||
|
||||
std::vector<common_peg_ast_id> entries;
|
||||
minimax_m3_collect(arena, node, tag, entries);
|
||||
|
||||
std::string result = is_object ? "{" : "[";
|
||||
|
||||
bool add_comma = false;
|
||||
for (auto entry_id : entries) {
|
||||
const auto & entry = arena.get(entry_id);
|
||||
|
||||
std::string text;
|
||||
if (is_object) {
|
||||
text = minimax_m3_member_to_json(arena, entry);
|
||||
} else {
|
||||
text = minimax_m3_value_to_json(arena, minimax_m3_value_of(arena, entry), !entry.is_partial);
|
||||
}
|
||||
|
||||
if (text.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (add_comma) {
|
||||
result += ",";
|
||||
}
|
||||
add_comma = true;
|
||||
result += text;
|
||||
}
|
||||
|
||||
if (closed) {
|
||||
result += is_object ? "}" : "]";
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
static std::string minimax_m3_value_to_json(const common_peg_ast_arena & arena, common_peg_ast_id id, bool closed) {
|
||||
if (id == COMMON_PEG_INVALID_AST_ID) {
|
||||
return "";
|
||||
}
|
||||
|
||||
const auto & node = arena.get(id);
|
||||
|
||||
if (node.tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_OBJECT) {
|
||||
return minimax_m3_container_to_json(arena, node, /* is_object = */ true, closed);
|
||||
}
|
||||
|
||||
if (node.tag == common_chat_peg_minimax_m3_mapper::TOOL_ARG_ARRAY) {
|
||||
return minimax_m3_container_to_json(arena, node, /* is_object = */ false, closed);
|
||||
}
|
||||
|
||||
if (node.tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE) {
|
||||
return "\"" + escape_json_string_inner(std::string(node.text)) + (closed ? "\"" : "");
|
||||
}
|
||||
|
||||
// Numbers and booleans are written verbatim by the template
|
||||
return std::string(node.text);
|
||||
}
|
||||
|
||||
void common_chat_peg_minimax_m3_mapper::from_ast(const common_peg_ast_arena & arena,
|
||||
const common_peg_parse_result & result) {
|
||||
for (const auto & node : result.nodes) {
|
||||
visit(arena, node);
|
||||
}
|
||||
}
|
||||
|
||||
void common_chat_peg_minimax_m3_mapper::visit(const common_peg_ast_arena & arena, common_peg_ast_id id) {
|
||||
const auto & node = arena.get(id);
|
||||
|
||||
if (node.tag == common_chat_peg_builder::REASONING) {
|
||||
result.reasoning_content += std::string(node.text);
|
||||
return;
|
||||
}
|
||||
|
||||
if (node.tag == common_chat_peg_builder::CONTENT) {
|
||||
result.content += std::string(node.text);
|
||||
return;
|
||||
}
|
||||
|
||||
if (node.tag == common_chat_peg_builder::TOOL) {
|
||||
auto name_id = arena.find_by_tag(node, common_chat_peg_builder::TOOL_NAME);
|
||||
if (name_id != COMMON_PEG_INVALID_AST_ID) {
|
||||
common_chat_tool_call call;
|
||||
call.name = std::string(arena.get(name_id).text);
|
||||
call.arguments = minimax_m3_container_to_json(arena, node, /* is_object = */ true, !node.is_partial);
|
||||
result.tool_calls.push_back(call);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
for (auto child_id : node.children) {
|
||||
visit(arena, child_id);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -40,6 +40,18 @@ class common_chat_peg_gemma4_mapper : public common_chat_peg_mapper {
|
||||
void visit(const common_peg_ast_arena & arena, common_peg_ast_id id);
|
||||
};
|
||||
|
||||
class common_chat_peg_minimax_m3_mapper : public common_chat_peg_mapper {
|
||||
public:
|
||||
static constexpr const char * TOOL_ARG_OBJECT = "tool-arg-object";
|
||||
static constexpr const char * TOOL_ARG_ARRAY = "tool-arg-array";
|
||||
static constexpr const char * TOOL_ARG_ITEM = "tool-arg-item";
|
||||
|
||||
common_chat_peg_minimax_m3_mapper(common_chat_msg & msg) : common_chat_peg_mapper(msg) {}
|
||||
virtual void from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result);
|
||||
private:
|
||||
void visit(const common_peg_ast_arena & arena, common_peg_ast_id id);
|
||||
};
|
||||
|
||||
struct content_structure;
|
||||
struct tool_call_structure;
|
||||
|
||||
|
||||
+301
-13
@@ -816,6 +816,8 @@ const char * common_chat_format_name(common_chat_format format) {
|
||||
return "peg-native";
|
||||
case COMMON_CHAT_FORMAT_PEG_GEMMA4:
|
||||
return "peg-gemma4";
|
||||
case COMMON_CHAT_FORMAT_PEG_MINIMAX_M3:
|
||||
return "peg-minimax-m3";
|
||||
default:
|
||||
throw std::runtime_error("Unknown chat format");
|
||||
}
|
||||
@@ -1024,7 +1026,7 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_
|
||||
|
||||
data.supports_thinking = true;
|
||||
data.thinking_start_tag = "[THINK]";
|
||||
data.thinking_end_tag = "[/THINK]";
|
||||
data.thinking_end_tags = {"[/THINK]"};
|
||||
data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages);
|
||||
data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override = */ adjusted_messages);
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
|
||||
@@ -1150,6 +1152,9 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
|
||||
data.supports_thinking = true;
|
||||
|
||||
data.thinking_start_tag = "<|channel|>analysis<|message|>";
|
||||
data.thinking_end_tags = {"<|end|>"};
|
||||
|
||||
// These special tokens are required to parse properly, so we include them
|
||||
// even if parse_tool_calls is false.
|
||||
data.preserved_tokens = {
|
||||
@@ -1294,7 +1299,7 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4;
|
||||
data.supports_thinking = true;
|
||||
data.thinking_start_tag = "<|channel>thought";
|
||||
data.thinking_end_tag = "<channel|>";
|
||||
data.thinking_end_tags = {"<channel|>"};
|
||||
|
||||
data.preserved_tokens = {
|
||||
"<|channel>",
|
||||
@@ -1569,7 +1574,7 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp
|
||||
const std::string GEN_PROMPT = "<|im_assistant|>assistant<|im_middle|>";
|
||||
|
||||
data.thinking_start_tag = THINK_START;
|
||||
data.thinking_end_tag = THINK_END;
|
||||
data.thinking_end_tags = {THINK_END};
|
||||
|
||||
if (inputs.has_continuation()) {
|
||||
const auto & msg = inputs.continue_msg;
|
||||
@@ -1703,7 +1708,7 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat
|
||||
}
|
||||
|
||||
data.thinking_start_tag = THINK_START;
|
||||
data.thinking_end_tag = THINK_END;
|
||||
data.thinking_end_tags = {THINK_END};
|
||||
|
||||
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
|
||||
auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
|
||||
@@ -1943,7 +1948,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
|
||||
data.supports_thinking = true;
|
||||
data.thinking_start_tag = "<think>";
|
||||
data.thinking_end_tag = "</think>";
|
||||
data.thinking_end_tags = {"</think>"};
|
||||
data.preserved_tokens = {
|
||||
"|DSML|",
|
||||
"<think>",
|
||||
@@ -2160,7 +2165,7 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
|
||||
data.supports_thinking = true;
|
||||
data.thinking_start_tag = THINK_START;
|
||||
data.thinking_end_tag = THINK_END;
|
||||
data.thinking_end_tags = {THINK_END};
|
||||
data.preserved_tokens = {
|
||||
TURN_START, TURN_END, CHATBOT, USER, SYSTEM,
|
||||
THINK_START, THINK_END,
|
||||
@@ -2179,9 +2184,10 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
|
||||
{ COMMON_CHAT_ROLE_SYSTEM, TURN_START + SYSTEM },
|
||||
};
|
||||
|
||||
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
|
||||
auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
|
||||
auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE;
|
||||
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
|
||||
auto has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty();
|
||||
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);
|
||||
|
||||
if (inputs.has_continuation()) {
|
||||
const auto & msg = inputs.continue_msg;
|
||||
@@ -2212,7 +2218,11 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
|
||||
p.optional(p.literal(THINK_END))));
|
||||
}
|
||||
|
||||
auto text_content = p.literal(TEXT_START) + p.content(p.until(TEXT_END)) + p.optional(p.literal(TEXT_END));
|
||||
auto text_content = has_response_format
|
||||
? p.literal(TEXT_START) +
|
||||
p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) +
|
||||
p.optional(p.literal(TEXT_END))
|
||||
: p.literal(TEXT_START) + p.content(p.until(TEXT_END)) + p.optional(p.literal(TEXT_END));
|
||||
|
||||
if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) {
|
||||
return generation_prompt + reasoning + text_content + p.optional(p.literal(TURN_END)) + end;
|
||||
@@ -2240,13 +2250,17 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
|
||||
data.parser = parser.save();
|
||||
|
||||
if (include_grammar) {
|
||||
data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
|
||||
data.grammar_lazy = !has_response_format && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
|
||||
data.grammar = build_grammar([&](const common_grammar_builder & builder) {
|
||||
foreach_function(inputs.tools, [&](const json & tool) {
|
||||
const auto & function = tool.at("function");
|
||||
auto schema = function.at("parameters");
|
||||
builder.resolve_refs(schema);
|
||||
});
|
||||
if (has_response_format) {
|
||||
auto schema = inputs.json_schema;
|
||||
builder.resolve_refs(schema);
|
||||
}
|
||||
parser.build_grammar(builder, data.grammar_lazy);
|
||||
});
|
||||
|
||||
@@ -2258,6 +2272,264 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
|
||||
return data;
|
||||
}
|
||||
|
||||
static common_chat_params common_chat_params_init_minimax_m3(const common_chat_template & tmpl,
|
||||
const autoparser::generation_params & inputs) {
|
||||
common_chat_params data;
|
||||
|
||||
data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs);
|
||||
data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs);
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_MINIMAX_M3;
|
||||
data.supports_thinking = true;
|
||||
data.thinking_start_tag = "<mm:think>";
|
||||
data.thinking_end_tags = {"</mm:think>"};
|
||||
|
||||
// M3 prefixes every tool tag with the namespace token "]<]minimax[>[";
|
||||
// params use the parameter name as the tag (<file_path>...</file_path>).
|
||||
const std::string NS = "]<]minimax[>[";
|
||||
const std::string THINK_START = "<mm:think>";
|
||||
const std::string THINK_END = "</mm:think>";
|
||||
const std::string FC_START = NS + "<tool_call>";
|
||||
const std::string FC_END = NS + "</tool_call>";
|
||||
const std::string INVOKE_END = NS + "</invoke>";
|
||||
|
||||
data.preserved_tokens = {
|
||||
NS,
|
||||
"<tool_call>",
|
||||
"</tool_call>",
|
||||
THINK_START,
|
||||
THINK_END,
|
||||
};
|
||||
|
||||
data.message_delimiters = {
|
||||
{ COMMON_CHAT_ROLE_ASSISTANT, "]~b]ai" },
|
||||
{ COMMON_CHAT_ROLE_USER, "]~b]user" },
|
||||
{ COMMON_CHAT_ROLE_TOOL, "]~b]tool" },
|
||||
{ COMMON_CHAT_ROLE_SYSTEM, "]~b]developer" },
|
||||
{ COMMON_CHAT_ROLE_SYSTEM, "]~b]system" },
|
||||
};
|
||||
|
||||
auto has_tools = inputs.tools.is_array() && !inputs.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);
|
||||
|
||||
const std::string GEN_PROMPT = data.generation_prompt;
|
||||
|
||||
using mm3 = common_chat_peg_minimax_m3_mapper;
|
||||
|
||||
if (inputs.has_continuation()) {
|
||||
const auto & msg = inputs.continue_msg;
|
||||
|
||||
data.generation_prompt = GEN_PROMPT + THINK_START + msg.reasoning_content;
|
||||
if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
|
||||
data.generation_prompt += THINK_END + msg.render_content();
|
||||
}
|
||||
|
||||
data.prompt += data.generation_prompt;
|
||||
}
|
||||
|
||||
auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
|
||||
auto generation_prompt = p.prefix(GEN_PROMPT, THINK_START);
|
||||
auto end = p.end();
|
||||
|
||||
auto reasoning = p.eps();
|
||||
if (extract_reasoning) {
|
||||
auto block = inputs.enable_thinking
|
||||
? p.literal(THINK_START) + p.space() +
|
||||
p.ac(p.reasoning(p.until(THINK_END)) + p.literal(THINK_END), THINK_END)
|
||||
: p.literal(THINK_START) + p.ac(p.until(THINK_END) + p.literal(THINK_END), THINK_END);
|
||||
|
||||
// A turn without reasoning is prefixed with a bare </mm:think>, written either by the
|
||||
// generation prompt (thinking_mode = "disabled") or by the model itself.
|
||||
reasoning = p.optional(p.choice({ block, p.literal(THINK_END) }));
|
||||
}
|
||||
|
||||
if (has_response_format) {
|
||||
auto response_format = p.rule("response-format",
|
||||
p.literal("```json") + p.space() +
|
||||
p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) +
|
||||
p.space() + p.literal("```"));
|
||||
return generation_prompt + reasoning + response_format + end;
|
||||
}
|
||||
|
||||
if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) {
|
||||
return generation_prompt + reasoning + p.content(p.rest()) + end;
|
||||
}
|
||||
|
||||
auto alternatives_of = [](const json & schema) -> std::optional<json> {
|
||||
for (const auto * keyword : { "oneOf", "anyOf" }) {
|
||||
if (schema.contains(keyword) && schema.at(keyword).is_array() && !schema.at(keyword).empty()) {
|
||||
return schema.at(keyword);
|
||||
}
|
||||
}
|
||||
return std::nullopt;
|
||||
};
|
||||
|
||||
auto tool_choice = p.choice();
|
||||
foreach_function(inputs.tools, [&](const json & tool) {
|
||||
const auto & function = tool.at("function");
|
||||
std::string name = function.at("name");
|
||||
auto params = function.contains("parameters") ? function.at("parameters") : json::object();
|
||||
|
||||
auto schema_info = common_schema_info();
|
||||
schema_info.resolve_refs(params);
|
||||
|
||||
// The template expands argument values recursively in XML (see the to_xml() macro)
|
||||
std::function<common_peg_parser(const json &, const std::string &, const std::string &)> value_of;
|
||||
std::function<common_peg_parser(const json &, const std::string &)> members_of;
|
||||
|
||||
auto element_of = [&](const std::string & tag, const json & schema, const std::string & rule_name) {
|
||||
const std::string close = NS + "</" + tag + ">";
|
||||
return p.rule(rule_name,
|
||||
p.tool_arg(
|
||||
p.tool_arg_open(
|
||||
p.literal(NS + "<") +
|
||||
p.tool_arg_name(p.literal(tag)) +
|
||||
p.literal(">")) +
|
||||
value_of(schema, rule_name, close)));
|
||||
};
|
||||
|
||||
value_of = [&](const json & schema,
|
||||
const std::string & rule_name,
|
||||
const std::string & close) -> common_peg_parser {
|
||||
auto close_tag = p.tool_arg_close(p.literal(close));
|
||||
|
||||
// A string accepts anything, so a union with a string alternative is a string
|
||||
if (schema_info.resolves_to_string(schema)) {
|
||||
return p.ac(p.tool_arg_string_value(p.until(close)) + close_tag, close);
|
||||
}
|
||||
|
||||
if (auto alternatives = alternatives_of(schema)) {
|
||||
std::vector<common_peg_parser> choices;
|
||||
|
||||
size_t index = 0;
|
||||
for (const auto & alternative : *alternatives) {
|
||||
const std::string alt_name = rule_name + "-" + std::to_string(index++);
|
||||
|
||||
// There is a risk that this breaks streaming deltas, but that's a risk we
|
||||
// assume to provide tool arg streaming.
|
||||
choices.push_back(value_of(alternative, alt_name, close));
|
||||
}
|
||||
|
||||
return p.choice(choices);
|
||||
}
|
||||
|
||||
const std::string type = schema.contains("type") && schema.at("type").is_string()
|
||||
? schema.at("type").get<std::string>()
|
||||
: "";
|
||||
|
||||
if (type == "object" && schema.contains("properties")) {
|
||||
return p.tag(mm3::TOOL_ARG_OBJECT, members_of(schema, rule_name)) + p.space() + close_tag;
|
||||
}
|
||||
|
||||
if (type == "array" && schema.contains("items")) {
|
||||
const std::string item_close = NS + "</item>";
|
||||
auto item = p.rule(rule_name + "-item",
|
||||
p.tag(mm3::TOOL_ARG_ITEM,
|
||||
p.literal(NS + "<item>") +
|
||||
value_of(schema.at("items"), rule_name + "-item", item_close)));
|
||||
return p.tag(mm3::TOOL_ARG_ARRAY, p.repeat(p.space() + item, 0, -1)) + p.space() + close_tag;
|
||||
}
|
||||
|
||||
return p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", schema, false)) + close_tag;
|
||||
};
|
||||
|
||||
// Required properties in schema order, then any number of optional ones in any order.
|
||||
members_of = [&](const json & schema, const std::string & rule_prefix) -> common_peg_parser {
|
||||
const auto & props = schema.at("properties");
|
||||
|
||||
std::set<std::string> required;
|
||||
if (schema.contains("required")) {
|
||||
schema.at("required").get_to(required);
|
||||
}
|
||||
|
||||
std::vector<common_peg_parser> required_elements;
|
||||
std::vector<common_peg_parser> optional_elements;
|
||||
for (const auto & [key, key_schema] : props.items()) {
|
||||
auto element = element_of(key, key_schema, rule_prefix + "-" + key);
|
||||
if (required.find(key) != required.end()) {
|
||||
required_elements.push_back(element);
|
||||
} else {
|
||||
optional_elements.push_back(element);
|
||||
}
|
||||
}
|
||||
|
||||
common_peg_parser members = p.eps();
|
||||
for (size_t i = 0; i < required_elements.size(); i++) {
|
||||
if (i > 0) {
|
||||
members = members + p.space();
|
||||
}
|
||||
members = members + required_elements[i];
|
||||
}
|
||||
|
||||
if (!optional_elements.empty()) {
|
||||
common_peg_parser any_optional = p.choice();
|
||||
for (const auto & element : optional_elements) {
|
||||
any_optional |= element;
|
||||
}
|
||||
members = members + p.repeat(p.space() + any_optional, 0, -1);
|
||||
}
|
||||
|
||||
return members;
|
||||
};
|
||||
|
||||
common_peg_parser invoke_body =
|
||||
params.contains("properties") ? members_of(params, "tool-" + name + "-arg") : p.eps();
|
||||
|
||||
auto func_parser = p.tool(
|
||||
p.tool_open(p.literal(NS + "<invoke name=\"") +
|
||||
p.tool_name(p.literal(name)) + p.literal("\">")) +
|
||||
p.space() + invoke_body + p.space() +
|
||||
p.tool_close(p.literal(INVOKE_END)));
|
||||
|
||||
tool_choice |= p.rule("tool-" + name, func_parser);
|
||||
});
|
||||
|
||||
auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;
|
||||
|
||||
common_peg_parser tool_calls = p.eps();
|
||||
if (inputs.parallel_tool_calls) {
|
||||
tool_calls = p.trigger_rule("tool-call",
|
||||
p.literal(FC_START) + p.space() + tool_choice +
|
||||
p.zero_or_more(p.space() + tool_choice) + p.space() + p.literal(FC_END));
|
||||
} else {
|
||||
tool_calls = p.trigger_rule("tool-call",
|
||||
p.literal(FC_START) + p.space() + tool_choice + p.space() + p.literal(FC_END));
|
||||
}
|
||||
|
||||
if (!require_tools) {
|
||||
tool_calls = p.optional(tool_calls);
|
||||
}
|
||||
|
||||
auto content_before_tools = p.content(p.until(FC_START));
|
||||
return generation_prompt + reasoning + content_before_tools + tool_calls + end;
|
||||
});
|
||||
|
||||
data.parser = parser.save();
|
||||
|
||||
if (include_grammar) {
|
||||
data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED));
|
||||
data.grammar = build_grammar([&](const common_grammar_builder & builder) {
|
||||
foreach_function(inputs.tools, [&](const json & tool) {
|
||||
const auto & function = tool.at("function");
|
||||
auto schema = function.contains("parameters") ? function.at("parameters") : json::object();
|
||||
builder.resolve_refs(schema);
|
||||
});
|
||||
if (has_response_format) {
|
||||
auto schema = inputs.json_schema;
|
||||
builder.resolve_refs(schema);
|
||||
}
|
||||
parser.build_grammar(builder, data.grammar_lazy);
|
||||
});
|
||||
|
||||
data.grammar_triggers = {
|
||||
{ COMMON_GRAMMAR_TRIGGER_TYPE_WORD, FC_START },
|
||||
};
|
||||
}
|
||||
|
||||
return data;
|
||||
}
|
||||
|
||||
namespace workaround {
|
||||
|
||||
static void map_developer_role_to_system(json & messages) {
|
||||
@@ -2501,7 +2773,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem
|
||||
};
|
||||
|
||||
data.thinking_start_tag = "<think>";
|
||||
data.thinking_end_tag = "</think>";
|
||||
data.thinking_end_tags = {"</think>"};
|
||||
|
||||
data.message_delimiters = {
|
||||
{ COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" },
|
||||
@@ -2695,6 +2967,15 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
|
||||
return common_chat_params_init_gigachat_v3(tmpl, params);
|
||||
}
|
||||
|
||||
// MiniMax-M3: the namespace token "]<]minimax[>[" collides with the autoparser's
|
||||
// markup delimiters, so detect the template and use a dedicated parser.
|
||||
if (src.find("]<]minimax[>[") != std::string::npos &&
|
||||
src.find("<tool_call>") != std::string::npos &&
|
||||
src.find("<invoke name=") != std::string::npos) {
|
||||
LOG_DBG("Using specialized template: MiniMax-M3\n");
|
||||
return common_chat_params_init_minimax_m3(tmpl, params);
|
||||
}
|
||||
|
||||
// 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".
|
||||
@@ -2857,7 +3138,10 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_
|
||||
auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE;
|
||||
if (auto_params.supports_thinking) {
|
||||
auto_params.thinking_start_tag = trim_whitespace(autoparser.reasoning.start);
|
||||
auto_params.thinking_end_tag = trim_whitespace(autoparser.reasoning.end);
|
||||
auto end_tag = trim_whitespace(autoparser.reasoning.end);
|
||||
if (!end_tag.empty()) {
|
||||
auto_params.thinking_end_tags = {std::move(end_tag)};
|
||||
}
|
||||
}
|
||||
common_peg_arena arena;
|
||||
arena.load(auto_params.parser);
|
||||
@@ -2983,6 +3267,8 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena & src_pars
|
||||
std::unique_ptr<common_chat_peg_mapper> mapper;
|
||||
if (params.format == COMMON_CHAT_FORMAT_PEG_GEMMA4) {
|
||||
mapper = std::make_unique<common_chat_peg_gemma4_mapper>(msg);
|
||||
} else if (params.format == COMMON_CHAT_FORMAT_PEG_MINIMAX_M3) {
|
||||
mapper = std::make_unique<common_chat_peg_minimax_m3_mapper>(msg);
|
||||
} else {
|
||||
mapper = std::make_unique<common_chat_peg_mapper>(msg);
|
||||
}
|
||||
@@ -3005,6 +3291,8 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena & src_pars
|
||||
std::unique_ptr<common_chat_peg_mapper> mapper;
|
||||
if (params.format == COMMON_CHAT_FORMAT_PEG_GEMMA4) {
|
||||
mapper = std::make_unique<common_chat_peg_gemma4_mapper>(msg);
|
||||
} else if (params.format == COMMON_CHAT_FORMAT_PEG_MINIMAX_M3) {
|
||||
mapper = std::make_unique<common_chat_peg_minimax_m3_mapper>(msg);
|
||||
} else {
|
||||
mapper = std::make_unique<common_chat_peg_mapper>(msg);
|
||||
}
|
||||
|
||||
+2
-1
@@ -233,6 +233,7 @@ enum common_chat_format {
|
||||
COMMON_CHAT_FORMAT_PEG_SIMPLE,
|
||||
COMMON_CHAT_FORMAT_PEG_NATIVE,
|
||||
COMMON_CHAT_FORMAT_PEG_GEMMA4,
|
||||
COMMON_CHAT_FORMAT_PEG_MINIMAX_M3,
|
||||
|
||||
COMMON_CHAT_FORMAT_COUNT, // Not a format, just the # formats
|
||||
};
|
||||
@@ -274,7 +275,7 @@ struct common_chat_params {
|
||||
std::string generation_prompt;
|
||||
bool supports_thinking = false;
|
||||
std::string thinking_start_tag; // e.g., "<think>"
|
||||
std::string thinking_end_tag; // e.g., "</think>"
|
||||
std::vector<std::string> thinking_end_tags; // e.g., "</think>"
|
||||
std::vector<common_grammar_trigger> grammar_triggers;
|
||||
std::vector<std::string> preserved_tokens;
|
||||
std::vector<std::string> additional_stops;
|
||||
|
||||
+29
-4
@@ -1249,7 +1249,6 @@ common_init_result::common_init_result(common_params & params, bool model_only)
|
||||
lora.reset(llama_adapter_lora_init(model, la.path.c_str()));
|
||||
if (lora == nullptr) {
|
||||
COM_ERR("failed to load lora adapter '%s'\n", la.path.c_str());
|
||||
pimpl->model.reset(model);
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -1519,23 +1518,49 @@ done:
|
||||
return res;
|
||||
}
|
||||
|
||||
void common_context_seq_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
|
||||
static void common_context_seq_rm(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
|
||||
auto * mem = llama_get_memory(ctx);
|
||||
if (!llama_memory_seq_rm(mem, seq_id, p0, p1)) {
|
||||
GGML_ABORT("%s", string_format("failed to remove sequence %d with p0=%d, p1=%d\n", seq_id, p0, p1).c_str());
|
||||
}
|
||||
}
|
||||
|
||||
void common_context_seq_cp(llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
|
||||
static void common_context_seq_cp(llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
|
||||
auto * mem = llama_get_memory(ctx);
|
||||
llama_memory_seq_cp(mem, seq_id_src, seq_id_dst, p0, p1);
|
||||
}
|
||||
|
||||
void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) {
|
||||
static void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) {
|
||||
auto * mem = llama_get_memory(ctx);
|
||||
llama_memory_seq_add(mem, seq_id, p0, p1, delta);
|
||||
}
|
||||
|
||||
void common_memory::init(llama_context * ctx_tgt, llama_context * ctx_dft) {
|
||||
this->ctx_tgt = ctx_tgt;
|
||||
this->ctx_dft = ctx_dft;
|
||||
}
|
||||
|
||||
void common_memory::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) const {
|
||||
common_context_seq_rm(ctx_tgt, seq_id, p0, p1);
|
||||
if (ctx_dft) {
|
||||
common_context_seq_rm(ctx_dft, seq_id, p0, p1);
|
||||
}
|
||||
}
|
||||
|
||||
void common_memory::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) const {
|
||||
common_context_seq_cp(ctx_tgt, seq_id_src, seq_id_dst, p0, p1);
|
||||
if (ctx_dft) {
|
||||
common_context_seq_cp(ctx_dft, seq_id_src, seq_id_dst, p0, p1);
|
||||
}
|
||||
}
|
||||
|
||||
void common_memory::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) const {
|
||||
common_context_seq_add(ctx_tgt, seq_id, p0, p1, delta);
|
||||
if (ctx_dft) {
|
||||
common_context_seq_add(ctx_dft, seq_id, p0, p1, delta);
|
||||
}
|
||||
}
|
||||
|
||||
void common_set_adapter_lora(struct llama_context * ctx, std::vector<common_adapter_lora_info> & lora) {
|
||||
std::vector<llama_adapter_lora *> loras;
|
||||
std::vector<float> scales;
|
||||
|
||||
+23
-11
@@ -173,6 +173,7 @@ enum common_speculative_type {
|
||||
COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, // Eagle3 speculative decoding
|
||||
COMMON_SPECULATIVE_TYPE_DRAFT_MTP, // Multi-token prediction
|
||||
COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, // DFlash speculative decoding
|
||||
COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, // DSpark speculative decoding (DFlash + Markov head)
|
||||
COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding based on n-grams
|
||||
COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only
|
||||
COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values
|
||||
@@ -284,12 +285,12 @@ struct common_params_sampling {
|
||||
|
||||
// reasoning budget sampler parameters
|
||||
// these are populated by the server/CLI based on chat template params
|
||||
int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget
|
||||
std::vector<llama_token> reasoning_budget_start; // start tag token sequence
|
||||
std::vector<llama_token> reasoning_budget_end; // end tag token sequence
|
||||
std::vector<llama_token> reasoning_budget_forced; // forced sequence (message + end tag)
|
||||
std::string reasoning_budget_message; // message injected before end tag when budget exhausted
|
||||
bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime
|
||||
int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget
|
||||
std::vector<llama_token> reasoning_budget_start; // start tag token sequence
|
||||
std::vector<llama_tokens> reasoning_budget_end; // end tag token sequences; the first tag is used as the forcing sequence
|
||||
std::vector<llama_token> reasoning_budget_forced; // forced sequence (message + first end tag)
|
||||
std::string reasoning_budget_message; // message injected before end tag when budget exhausted
|
||||
bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime
|
||||
|
||||
bool backend_sampling = false;
|
||||
|
||||
@@ -388,7 +389,7 @@ struct common_params_speculative {
|
||||
|
||||
uint32_t need_n_rs_seq() const {
|
||||
bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) {
|
||||
return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH;
|
||||
return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
|
||||
});
|
||||
|
||||
return needs_rs_seq ? draft.n_max : 0u;
|
||||
@@ -668,6 +669,10 @@ struct common_params {
|
||||
// enable built-in tools
|
||||
std::vector<std::string> server_tools;
|
||||
|
||||
// MCP server configs (Cursor-compatible JSON)
|
||||
std::string mcp_servers_config; // path to JSON file with MCP server definitions
|
||||
std::string mcp_servers_json; // inline JSON with MCP server definitions
|
||||
|
||||
// router server configs
|
||||
std::string models_dir = ""; // directory containing models for the router server
|
||||
std::string models_preset = ""; // directory containing model presets for the router server
|
||||
@@ -944,10 +949,17 @@ enum common_context_seq_rm_type {
|
||||
// note: clears the memory of the context
|
||||
common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx);
|
||||
|
||||
// aborts execution on failure
|
||||
void common_context_seq_rm (llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1);
|
||||
void common_context_seq_add(llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta);
|
||||
void common_context_seq_cp (llama_context * ctx, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1);
|
||||
struct common_memory {
|
||||
llama_context * ctx_tgt = nullptr;
|
||||
llama_context * ctx_dft = nullptr;
|
||||
|
||||
void init(llama_context * ctx_tgt, llama_context * ctx_dft = nullptr);
|
||||
|
||||
// aborts execution on failure
|
||||
void seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) const;
|
||||
void seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) const;
|
||||
void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) const;
|
||||
};
|
||||
|
||||
//
|
||||
// Batch utils
|
||||
|
||||
+52
-17
@@ -568,16 +568,30 @@ static hf_cache::hf_files get_split_files(const hf_cache::hf_files & files,
|
||||
}
|
||||
|
||||
// pick the best sibling GGUF whose filename contains `keyword` (e.g. "mmproj" / "mtp"),
|
||||
// preferring deeper shared directory prefix with the model, then closest quantization
|
||||
// preferring deeper shared directory prefix with the model, then exact `tag` match,
|
||||
// then closest quantization to the tag when given, or to the model otherwise
|
||||
static hf_cache::hf_file find_best_sibling(const hf_cache::hf_files & files,
|
||||
const std::string & model,
|
||||
const std::string & keyword) {
|
||||
const std::string & keyword,
|
||||
const std::string & tag = "") {
|
||||
hf_cache::hf_file best;
|
||||
size_t best_depth = 0;
|
||||
int best_diff = 0;
|
||||
bool best_exact = false;
|
||||
bool found = false;
|
||||
|
||||
auto model_bits = extract_quant_bits(model);
|
||||
std::string tag_upper = tag;
|
||||
for (char & c : tag_upper) {
|
||||
c = (char) std::toupper((unsigned char) c);
|
||||
}
|
||||
|
||||
int model_bits = 0;
|
||||
if (!tag_upper.empty()) {
|
||||
auto pos = tag_upper.find_first_of("0123456789");
|
||||
model_bits = pos == std::string::npos ? 0 : std::stoi(tag_upper.substr(pos));
|
||||
} else {
|
||||
model_bits = extract_quant_bits(model);
|
||||
}
|
||||
auto model_parts = string_split<std::string>(model, '/');
|
||||
auto model_dir = model_parts.end() - 1;
|
||||
|
||||
@@ -600,10 +614,19 @@ static hf_cache::hf_file find_best_sibling(const hf_cache::hf_files & files,
|
||||
auto bits = extract_quant_bits(f.path);
|
||||
auto diff = std::abs(bits - model_bits);
|
||||
|
||||
if (!found || depth > best_depth || (depth == best_depth && diff < best_diff)) {
|
||||
std::string path_upper = f.path;
|
||||
for (char & c : path_upper) {
|
||||
c = (char) std::toupper((unsigned char) c);
|
||||
}
|
||||
bool exact = !tag_upper.empty() && path_upper.find("-" + tag_upper + ".") != std::string::npos;
|
||||
|
||||
if (!found || depth > best_depth ||
|
||||
(depth == best_depth && exact && !best_exact) ||
|
||||
(depth == best_depth && exact == best_exact && diff < best_diff)) {
|
||||
best = f;
|
||||
best_depth = depth;
|
||||
best_diff = diff;
|
||||
best_exact = exact;
|
||||
found = true;
|
||||
}
|
||||
}
|
||||
@@ -616,18 +639,21 @@ static hf_cache::hf_file find_best_mmproj(const hf_cache::hf_files & files,
|
||||
}
|
||||
|
||||
static hf_cache::hf_file find_best_mtp(const hf_cache::hf_files & files,
|
||||
const std::string & model) {
|
||||
return find_best_sibling(files, model, "mtp-");
|
||||
const std::string & model,
|
||||
const std::string & tag = "") {
|
||||
return find_best_sibling(files, model, "mtp-", tag);
|
||||
}
|
||||
|
||||
static hf_cache::hf_file find_best_eagle3(const hf_cache::hf_files & files,
|
||||
const std::string & model) {
|
||||
return find_best_sibling(files, model, "eagle3-");
|
||||
const std::string & model,
|
||||
const std::string & tag = "") {
|
||||
return find_best_sibling(files, model, "eagle3-", tag);
|
||||
}
|
||||
|
||||
static hf_cache::hf_file find_best_dflash(const hf_cache::hf_files & files,
|
||||
const std::string & model) {
|
||||
return find_best_sibling(files, model, "dflash-");
|
||||
const std::string & model,
|
||||
const std::string & tag = "") {
|
||||
return find_best_sibling(files, model, "dflash-", tag);
|
||||
}
|
||||
|
||||
static bool gguf_filename_is_model(const std::string & filepath) {
|
||||
@@ -736,27 +762,36 @@ common_download_hf_plan common_download_get_hf_plan(const common_params_model &
|
||||
}
|
||||
} else {
|
||||
primary = find_best_model(all, tag);
|
||||
if (primary.path.empty()) {
|
||||
// a requested sidecar can resolve on its own, without a full model of the same tag
|
||||
if (primary.path.empty() && !opts.download_mtp && !opts.download_dflash && !opts.download_eagle3) {
|
||||
LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str());
|
||||
list_available_gguf_files(all);
|
||||
return plan;
|
||||
}
|
||||
}
|
||||
|
||||
plan.primary = primary;
|
||||
plan.model_files = get_split_files(all, primary);
|
||||
if (!primary.path.empty()) {
|
||||
plan.primary = primary;
|
||||
plan.model_files = get_split_files(all, primary);
|
||||
}
|
||||
|
||||
if (opts.download_mmproj) {
|
||||
if (opts.download_mmproj && !primary.path.empty()) {
|
||||
plan.mmproj = find_best_mmproj(all, primary.path);
|
||||
}
|
||||
if (opts.download_mtp) {
|
||||
plan.mtp = find_best_mtp(all, primary.path);
|
||||
plan.mtp = find_best_mtp(all, primary.path, tag);
|
||||
}
|
||||
if (opts.download_dflash) {
|
||||
plan.dflash = find_best_dflash(all, primary.path);
|
||||
plan.dflash = find_best_dflash(all, primary.path, tag);
|
||||
}
|
||||
if (opts.download_eagle3) {
|
||||
plan.eagle3 = find_best_eagle3(all, primary.path);
|
||||
plan.eagle3 = find_best_eagle3(all, primary.path, tag);
|
||||
}
|
||||
|
||||
if (primary.path.empty() &&
|
||||
plan.mtp.local_path.empty() && plan.dflash.local_path.empty() && plan.eagle3.local_path.empty()) {
|
||||
LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str());
|
||||
list_available_gguf_files(all);
|
||||
}
|
||||
|
||||
return plan;
|
||||
|
||||
+1
-1
@@ -136,7 +136,7 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(
|
||||
devs.push_back(llama_model_get_device(model, i));
|
||||
}
|
||||
|
||||
hp_ngl = llama_model_n_layer(model);
|
||||
hp_ngl = llama_model_n_layer(model) + llama_model_n_layer_nextn(model);
|
||||
hp_n_ctx_train = llama_model_n_ctx_train(model);
|
||||
hp_n_expert = llama_model_n_expert(model);
|
||||
|
||||
|
||||
+5
-153
@@ -3,10 +3,10 @@
|
||||
#include "common.h"
|
||||
#include "json-schema-to-grammar.h"
|
||||
#include "log.h"
|
||||
#include "trie.h"
|
||||
#include "unicode.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <deque>
|
||||
#include <initializer_list>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
@@ -32,154 +32,6 @@ static bool is_hex_digit(const char c) {
|
||||
return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F');
|
||||
}
|
||||
|
||||
// Trie for matching multiple literals.
|
||||
// This is used in common_peg_until_parser and to build a GBNF exclusion grammar
|
||||
struct trie {
|
||||
struct node {
|
||||
std::map<uint32_t, size_t> children; // Use uint32_t to store Unicode codepoints
|
||||
bool is_word;
|
||||
};
|
||||
|
||||
std::vector<node> nodes;
|
||||
|
||||
trie(const std::vector<std::string> & words) {
|
||||
create_node(); // root node
|
||||
for (const auto & w : words) {
|
||||
insert(w);
|
||||
}
|
||||
}
|
||||
|
||||
enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH };
|
||||
|
||||
// Check if a delimiter starts at the given position
|
||||
match_result check_at(std::string_view sv, size_t start_pos) const {
|
||||
size_t current = 0; // Start at root
|
||||
size_t pos = start_pos;
|
||||
|
||||
// LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str());
|
||||
|
||||
while (pos < sv.size()) {
|
||||
auto result = common_parse_utf8_codepoint(sv, pos);
|
||||
if (result.status != utf8_parse_result::SUCCESS) {
|
||||
break;
|
||||
}
|
||||
|
||||
auto it = nodes[current].children.find(result.codepoint);
|
||||
if (it == nodes[current].children.end()) {
|
||||
// Can't continue matching
|
||||
return match_result{match_result::NO_MATCH};
|
||||
}
|
||||
|
||||
current = it->second;
|
||||
pos += result.bytes_consumed;
|
||||
|
||||
// Check if we've matched a complete word
|
||||
if (nodes[current].is_word) {
|
||||
return match_result{match_result::COMPLETE_MATCH};
|
||||
}
|
||||
}
|
||||
|
||||
// Reached end of input while still in the trie (not at root)
|
||||
if (current != 0) {
|
||||
// We're in the middle of a potential match
|
||||
return match_result{match_result::PARTIAL_MATCH};
|
||||
}
|
||||
|
||||
// Reached end at root (no match)
|
||||
return match_result{match_result::NO_MATCH};
|
||||
}
|
||||
|
||||
private:
|
||||
size_t create_node() {
|
||||
size_t index = nodes.size();
|
||||
nodes.emplace_back();
|
||||
return index;
|
||||
}
|
||||
|
||||
void insert(const std::string & word) {
|
||||
size_t current = 0;
|
||||
size_t pos = 0;
|
||||
while (pos < word.length()) {
|
||||
auto result = common_parse_utf8_codepoint(word, pos);
|
||||
if (result.status != utf8_parse_result::SUCCESS) {
|
||||
break;
|
||||
}
|
||||
|
||||
uint32_t ch = result.codepoint;
|
||||
pos += result.bytes_consumed;
|
||||
|
||||
auto it = nodes[current].children.find(ch);
|
||||
if (it == nodes[current].children.end()) {
|
||||
size_t child = create_node();
|
||||
nodes[current].children[ch] = child;
|
||||
current = child;
|
||||
} else {
|
||||
current = it->second;
|
||||
}
|
||||
}
|
||||
nodes[current].is_word = true;
|
||||
}
|
||||
};
|
||||
|
||||
// Aho-Corasick automaton
|
||||
struct aho_corasick {
|
||||
trie t;
|
||||
std::vector<size_t> fail; // failure links
|
||||
std::vector<size_t> order; // states in BFS order
|
||||
std::vector<bool> terminal; // match states (directly or via a suffix link)
|
||||
std::set<uint32_t> alphabet; // every character with a transition
|
||||
|
||||
aho_corasick(const std::vector<std::string> & strings) : t(strings) {
|
||||
const auto & nodes = t.nodes;
|
||||
const size_t n = nodes.size();
|
||||
|
||||
fail.assign(n, 0);
|
||||
order.reserve(n);
|
||||
|
||||
std::deque<size_t> queue{ 0 };
|
||||
while (!queue.empty()) {
|
||||
size_t u = queue.front();
|
||||
queue.pop_front();
|
||||
order.push_back(u);
|
||||
for (const auto & [ch, v] : nodes[u].children) {
|
||||
if (u != 0) {
|
||||
size_t f = fail[u];
|
||||
while (f && nodes[f].children.find(ch) == nodes[f].children.end()) {
|
||||
f = fail[f];
|
||||
}
|
||||
auto it = nodes[f].children.find(ch);
|
||||
fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0;
|
||||
}
|
||||
queue.push_back(v);
|
||||
}
|
||||
}
|
||||
|
||||
terminal.assign(n, false);
|
||||
for (size_t u : order) {
|
||||
terminal[u] = nodes[u].is_word || (u != 0 && terminal[fail[u]]);
|
||||
}
|
||||
|
||||
for (const auto & node : nodes) {
|
||||
for (const auto & [ch, v] : node.children) {
|
||||
alphabet.insert(ch);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
size_t num_states() const { return t.nodes.size(); }
|
||||
bool is_terminal(size_t s) const { return terminal[s]; }
|
||||
|
||||
// follow failure links until a transition on `ch` exists.
|
||||
size_t next(size_t state, uint32_t ch) const {
|
||||
const auto & nodes = t.nodes;
|
||||
while (state && nodes[state].children.find(ch) == nodes[state].children.end()) {
|
||||
state = fail[state];
|
||||
}
|
||||
auto it = nodes[state].children.find(ch);
|
||||
return it != nodes[state].children.end() ? it->second : 0;
|
||||
}
|
||||
};
|
||||
|
||||
static std::pair<uint32_t, size_t> parse_hex_escape(const std::string & str, size_t pos, int hex_count) {
|
||||
if (pos + hex_count > str.length()) {
|
||||
return {0, 0};
|
||||
@@ -797,7 +649,7 @@ struct parser_executor {
|
||||
}
|
||||
|
||||
common_peg_parse_result operator()(const common_peg_until_parser & p) const {
|
||||
trie matcher(p.delimiters);
|
||||
common_trie matcher(p.delimiters);
|
||||
|
||||
// Scan input and check for delimiters
|
||||
size_t pos = start_pos;
|
||||
@@ -824,12 +676,12 @@ struct parser_executor {
|
||||
// Check if a delimiter starts at this position
|
||||
auto match = matcher.check_at(ctx.input, pos);
|
||||
|
||||
if (match == trie::COMPLETE_MATCH) {
|
||||
if (match == common_trie::COMPLETE_MATCH) {
|
||||
// Found a complete delimiter, return everything before it
|
||||
return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos);
|
||||
}
|
||||
|
||||
if (match == trie::PARTIAL_MATCH) {
|
||||
if (match == common_trie::PARTIAL_MATCH) {
|
||||
// Found a partial match extending to end of input, return everything before it
|
||||
return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos);
|
||||
}
|
||||
@@ -1559,7 +1411,7 @@ static std::string gbnf_ac_grammar(
|
||||
const std::map<size_t, std::vector<uint32_t>> &,
|
||||
const std::vector<uint32_t> &,
|
||||
const std::function<std::string(size_t)> &)> & build_rule) {
|
||||
aho_corasick ac(strings);
|
||||
common_aho_corasick ac(strings);
|
||||
|
||||
auto state_name = [&](size_t s) -> std::string {
|
||||
if (s == 0) {
|
||||
|
||||
@@ -330,6 +330,10 @@ common_presets common_preset_context::load_from_ini(const std::string & path, co
|
||||
}
|
||||
}
|
||||
|
||||
if (preset.name == COMMON_PRESET_DEFAULT_NAME && preset.options.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (preset.name == "*") {
|
||||
// handle global preset
|
||||
global = preset;
|
||||
|
||||
+77
-39
@@ -1,39 +1,52 @@
|
||||
#include "reasoning-budget.h"
|
||||
#include "common.h"
|
||||
#include "trie.h"
|
||||
#include "unicode.h"
|
||||
|
||||
#include "log.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
struct token_matcher {
|
||||
std::vector<llama_token> tokens;
|
||||
size_t pos = 0;
|
||||
std::vector<llama_tokens> seqs;
|
||||
common_aho_corasick ac;
|
||||
size_t state = 0;
|
||||
|
||||
bool advance(llama_token token) {
|
||||
if (tokens.empty()) {
|
||||
return false;
|
||||
}
|
||||
token_matcher(const std::vector<llama_tokens> & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {}
|
||||
|
||||
if (token == tokens[pos]) {
|
||||
pos++;
|
||||
if (pos >= tokens.size()) {
|
||||
pos = 0;
|
||||
return true;
|
||||
}
|
||||
} else {
|
||||
pos = 0;
|
||||
if (token == tokens[0]) {
|
||||
pos = 1;
|
||||
static std::vector<llama_tokens> collect(const std::vector<llama_tokens> & seqs) {
|
||||
std::vector<llama_tokens> res;
|
||||
for (const auto & seq : seqs) {
|
||||
if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) {
|
||||
res.push_back(seq);
|
||||
}
|
||||
}
|
||||
return false;
|
||||
return res;
|
||||
}
|
||||
|
||||
void reset() { pos = 0; }
|
||||
static common_trie build_trie(const std::vector<llama_tokens> & seqs) {
|
||||
common_trie t;
|
||||
for (const auto & seq : seqs) {
|
||||
t.insert(std::vector<uint32_t>(seq.begin(), seq.end()));
|
||||
}
|
||||
return t;
|
||||
}
|
||||
|
||||
// returns the index into seqs of the longest sequence ending at this token, or -1
|
||||
int32_t advance(llama_token token) {
|
||||
state = ac.next(state, (uint32_t) token);
|
||||
const int32_t p = ac.match_pattern(state);
|
||||
if (p >= 0) {
|
||||
state = 0;
|
||||
}
|
||||
return p;
|
||||
}
|
||||
|
||||
void reset() { state = 0; }
|
||||
};
|
||||
|
||||
struct common_reasoning_budget_ctx {
|
||||
@@ -41,7 +54,7 @@ struct common_reasoning_budget_ctx {
|
||||
|
||||
token_matcher start_matcher;
|
||||
token_matcher end_matcher;
|
||||
std::vector<llama_token> forced_tokens;
|
||||
llama_tokens forced_tokens;
|
||||
|
||||
int32_t budget; // maximum tokens in reasoning block
|
||||
int32_t remaining; // tokens remaining in budget
|
||||
@@ -50,6 +63,8 @@ struct common_reasoning_budget_ctx {
|
||||
|
||||
// for forcing
|
||||
size_t force_pos; // next position in forced_tokens to force
|
||||
|
||||
int32_t end_match; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none
|
||||
};
|
||||
|
||||
static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) {
|
||||
@@ -62,7 +77,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
|
||||
switch (ctx->state) {
|
||||
case REASONING_BUDGET_IDLE:
|
||||
{
|
||||
if (ctx->start_matcher.advance(token)) {
|
||||
if (ctx->start_matcher.advance(token) >= 0) {
|
||||
ctx->state = REASONING_BUDGET_COUNTING;
|
||||
ctx->remaining = ctx->budget;
|
||||
COM_TRC("activated, budget=%d tokens\n", ctx->budget);
|
||||
@@ -78,8 +93,10 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
|
||||
case REASONING_BUDGET_COUNTING:
|
||||
case REASONING_BUDGET_WAITING_UTF8:
|
||||
{
|
||||
if (ctx->end_matcher.advance(token)) {
|
||||
const int32_t match = ctx->end_matcher.advance(token);
|
||||
if (match >= 0) {
|
||||
ctx->state = REASONING_BUDGET_DONE;
|
||||
ctx->end_match = match;
|
||||
COM_TRC("%s", "deactivated (natural end)\n");
|
||||
break;
|
||||
}
|
||||
@@ -115,19 +132,25 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
|
||||
break;
|
||||
}
|
||||
case REASONING_BUDGET_FORCING:
|
||||
{
|
||||
// track the end sequence within forced_tokens so it is also reported on DONE
|
||||
const int32_t match = ctx->end_matcher.advance(token);
|
||||
ctx->force_pos++;
|
||||
if (ctx->force_pos >= ctx->forced_tokens.size()) {
|
||||
ctx->state = REASONING_BUDGET_DONE;
|
||||
ctx->end_match = match;
|
||||
COM_TRC("%s", "forced sequence complete, done\n");
|
||||
}
|
||||
break;
|
||||
}
|
||||
case REASONING_BUDGET_DONE:
|
||||
// Re-arm on a new start tag: some models emit multiple <think> blocks
|
||||
// per response, and each should get a fresh budget window.
|
||||
if (ctx->start_matcher.advance(token)) {
|
||||
if (ctx->start_matcher.advance(token) >= 0) {
|
||||
ctx->state = REASONING_BUDGET_COUNTING;
|
||||
ctx->remaining = ctx->budget;
|
||||
ctx->end_matcher.reset();
|
||||
ctx->end_match = -1;
|
||||
COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget);
|
||||
|
||||
if (ctx->remaining <= 0) {
|
||||
@@ -169,11 +192,12 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) {
|
||||
ctx->start_matcher.reset();
|
||||
ctx->end_matcher.reset();
|
||||
ctx->force_pos = 0;
|
||||
ctx->end_match = -1;
|
||||
}
|
||||
|
||||
static struct llama_sampler * common_reasoning_budget_init_state(
|
||||
const struct llama_vocab * vocab, const std::vector<llama_token> & start_tokens,
|
||||
const std::vector<llama_token> & end_tokens, const std::vector<llama_token> & forced_tokens,
|
||||
const struct llama_vocab * vocab, const std::vector<llama_tokens> & start_seqs,
|
||||
const std::vector<llama_tokens> & end_seqs, const llama_tokens & forced_tokens,
|
||||
int32_t budget, common_reasoning_budget_state initial_state);
|
||||
|
||||
static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl);
|
||||
@@ -205,12 +229,12 @@ static struct llama_sampler * common_reasoning_budget_clone(const struct llama_s
|
||||
}
|
||||
|
||||
static struct llama_sampler * common_reasoning_budget_init_state(
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_token> & start_tokens,
|
||||
const std::vector<llama_token> & end_tokens,
|
||||
const std::vector<llama_token> & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state) {
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_tokens> & start_seqs,
|
||||
const std::vector<llama_tokens> & end_seqs,
|
||||
const llama_tokens & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state) {
|
||||
// promote COUNTING with budget <= 0 to FORCING
|
||||
if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) {
|
||||
initial_state = REASONING_BUDGET_FORCING;
|
||||
@@ -220,25 +244,26 @@ static struct llama_sampler * common_reasoning_budget_init_state(
|
||||
/* .iface = */ &common_reasoning_budget_i,
|
||||
/* .ctx = */ new common_reasoning_budget_ctx {
|
||||
/* .vocab = */ vocab,
|
||||
/* .start_matcher = */ { start_tokens, 0 },
|
||||
/* .end_matcher = */ { end_tokens, 0 },
|
||||
/* .start_matcher = */ token_matcher(start_seqs),
|
||||
/* .end_matcher = */ token_matcher(end_seqs),
|
||||
/* .forced_tokens = */ forced_tokens,
|
||||
/* .budget = */ budget,
|
||||
/* .remaining = */ budget,
|
||||
/* .state = */ initial_state,
|
||||
/* .force_pos = */ 0,
|
||||
/* .end_match = */ -1,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
struct llama_sampler * common_reasoning_budget_init(
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_token> & start_tokens,
|
||||
const std::vector<llama_token> & end_tokens,
|
||||
const std::vector<llama_token> & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state) {
|
||||
return common_reasoning_budget_init_state(vocab, start_tokens, end_tokens, forced_tokens, budget, initial_state);
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_tokens> & start_seqs,
|
||||
const std::vector<llama_tokens> & end_seqs,
|
||||
const llama_tokens & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state) {
|
||||
return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state);
|
||||
}
|
||||
|
||||
common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) {
|
||||
@@ -248,6 +273,19 @@ common_reasoning_budget_state common_reasoning_budget_get_state(const struct lla
|
||||
return ((const common_reasoning_budget_ctx *)smpl->ctx)->state;
|
||||
}
|
||||
|
||||
const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl) {
|
||||
if (!smpl) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx;
|
||||
if (ctx->end_match < 0) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
return &ctx->end_matcher.seqs[ctx->end_match];
|
||||
}
|
||||
|
||||
bool common_reasoning_budget_force(struct llama_sampler * smpl) {
|
||||
if (!smpl) {
|
||||
return false;
|
||||
|
||||
+16
-10
@@ -2,6 +2,8 @@
|
||||
|
||||
#include "llama.h"
|
||||
|
||||
#include "common.h"
|
||||
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
|
||||
@@ -17,30 +19,34 @@ enum common_reasoning_budget_state {
|
||||
// reasoning block (e.g. between <think> and </think>).
|
||||
//
|
||||
// State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE
|
||||
// IDLE: passthrough, watching for start_tokens sequence
|
||||
// COUNTING: counting down remaining tokens, watching for natural end_tokens
|
||||
// IDLE: passthrough, watching for a start sequence
|
||||
// COUNTING: counting down remaining tokens, watching for a natural end sequence
|
||||
// WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence
|
||||
// FORCING: forces forced_tokens token-by-token (all other logits -> -inf)
|
||||
// DONE: passthrough forever
|
||||
//
|
||||
// Parameters:
|
||||
// vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr)
|
||||
// start_tokens - token sequence that activates counting
|
||||
// end_tokens - token sequence for natural deactivation
|
||||
// start_seqs - token sequences, any of which activates counting
|
||||
// end_seqs - token sequences, any of which naturally deactivates
|
||||
// forced_tokens - token sequence forced when budget expires
|
||||
// budget - max tokens allowed in the reasoning block
|
||||
// initial_state - initial state
|
||||
//
|
||||
struct llama_sampler * common_reasoning_budget_init(
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_token> & start_tokens,
|
||||
const std::vector<llama_token> & end_tokens,
|
||||
const std::vector<llama_token> & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE);
|
||||
const struct llama_vocab * vocab,
|
||||
const std::vector<llama_tokens> & start_seqs,
|
||||
const std::vector<llama_tokens> & end_seqs,
|
||||
const llama_tokens & forced_tokens,
|
||||
int32_t budget,
|
||||
common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE);
|
||||
|
||||
common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl);
|
||||
|
||||
// The end sequence that transitioned the sampler to DONE, or nullptr if none
|
||||
// was recorded. Cleared when a new start sequence re-arms the sampler.
|
||||
const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl);
|
||||
|
||||
// Manually transition the reasoning budget sampler into the FORCING state.
|
||||
// Returns true if the transition occurred.
|
||||
bool common_reasoning_budget_force(struct llama_sampler * smpl);
|
||||
|
||||
+12
-1
@@ -299,7 +299,7 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st
|
||||
if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0 || params.reasoning_control)) {
|
||||
rbudget = common_reasoning_budget_init(
|
||||
vocab,
|
||||
params.reasoning_budget_start,
|
||||
{params.reasoning_budget_start},
|
||||
params.reasoning_budget_end,
|
||||
params.reasoning_budget_forced,
|
||||
params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens);
|
||||
@@ -453,6 +453,17 @@ void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, boo
|
||||
|
||||
if (gsmpl->rbudget && is_generated) {
|
||||
llama_sampler_accept(gsmpl->rbudget, token);
|
||||
|
||||
// if done, replay end sequence which may contain a grammar trigger
|
||||
const bool is_done = common_reasoning_budget_get_state(gsmpl->rbudget) == REASONING_BUDGET_DONE;
|
||||
if (gsmpl->grmr && !accept_grammar && is_done) {
|
||||
const llama_tokens * end_seq = common_reasoning_budget_get_end_match(gsmpl->rbudget);
|
||||
if (end_seq) {
|
||||
for (const llama_token end_token : *end_seq) {
|
||||
llama_sampler_accept(gsmpl->grmr, end_token);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (gsmpl->grmr && accept_grammar) {
|
||||
|
||||
+91
-35
@@ -34,6 +34,7 @@ const std::map<std::string, common_speculative_type> common_speculative_type_fro
|
||||
{"draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3},
|
||||
{"draft-mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP},
|
||||
{"draft-dflash", COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH},
|
||||
{"draft-dspark", COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK},
|
||||
{"ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE},
|
||||
{"ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K},
|
||||
{"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V},
|
||||
@@ -437,6 +438,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
|
||||
int32_t n_embd_dec = 0; // draft hidden size
|
||||
int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size
|
||||
int32_t n_embd_tgt = 0; // target model hidden size
|
||||
int32_t n_layer_tgt = 0; // target model layer count
|
||||
|
||||
const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices
|
||||
uint32_t target_layer_ids_n = 0;
|
||||
@@ -478,6 +480,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
|
||||
n_embd_tgt = llama_model_n_embd(model_tgt);
|
||||
n_embd_dec = llama_model_n_embd(model_dft);
|
||||
n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt;
|
||||
n_layer_tgt = llama_model_n_layer(model_tgt);
|
||||
|
||||
const int32_t n_b = (int32_t) llama_n_batch(ctx_dft);
|
||||
batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd_dec, /*n_seq_max=*/ 1);
|
||||
@@ -510,9 +513,15 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
|
||||
}
|
||||
}
|
||||
|
||||
// turn on extraction of the target layers' input embeddings
|
||||
// turn on extraction of the target layers' hidden states
|
||||
for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
|
||||
llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
|
||||
if (target_layer_ids[k] < n_layer_tgt) {
|
||||
llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
|
||||
} else if (target_layer_ids[k] == n_layer_tgt) {
|
||||
llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false);
|
||||
} else {
|
||||
GGML_ABORT("EAGLE3: target layer id %d exceeds target n_layer %d", target_layer_ids[k], n_layer_tgt);
|
||||
}
|
||||
}
|
||||
|
||||
// turn on extraction of the draft model's pre-norm hidden state
|
||||
@@ -600,7 +609,9 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
|
||||
features_buf.resize((size_t) n_tokens * n_embd_enc, 0.0f);
|
||||
|
||||
for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
|
||||
const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
|
||||
const float * layer = target_layer_ids[k] < n_layer_tgt
|
||||
? llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k])
|
||||
: llama_get_embeddings_nextn(ctx_tgt);
|
||||
if (!layer) {
|
||||
GGML_ABORT("EAGLE3: target layer %d input not extracted.", target_layer_ids[k]);
|
||||
}
|
||||
@@ -918,15 +929,20 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
int32_t block_size = 0;
|
||||
llama_token mask_token_id = 0;
|
||||
|
||||
// draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
|
||||
const bool is_dspark;
|
||||
|
||||
const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices
|
||||
uint32_t target_layer_ids_n = 0;
|
||||
|
||||
// scratch buffer for concatenated target features [n_tokens, n_embd_enc]
|
||||
std::vector<float> features_buf;
|
||||
|
||||
common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq)
|
||||
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, n_seq)
|
||||
common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq,
|
||||
common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)
|
||||
: common_speculative_impl(type, n_seq)
|
||||
, params(params.draft)
|
||||
, is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)
|
||||
{
|
||||
auto * ctx_tgt = this->params.ctx_tgt;
|
||||
auto * ctx_dft = this->params.ctx_dft;
|
||||
@@ -953,16 +969,18 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
}
|
||||
mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));
|
||||
|
||||
LOG_INF("%s: adding speculative implementation 'draft-dflash'\n", __func__);
|
||||
LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str());
|
||||
LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);
|
||||
LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u\n", __func__, block_size, mask_token_id, target_layer_ids_n);
|
||||
|
||||
// DFlash input is [id_last, <mask> * (block_size-1)], so it can draft at most block_size-1 tokens per step
|
||||
if (this->params.n_max > block_size - 1 || this->params.n_min > block_size - 1) {
|
||||
LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained DFlash block size %d -- clamping to %d\n",
|
||||
__func__, this->params.n_max, this->params.n_min, block_size, block_size - 1);
|
||||
this->params.n_max = std::min(this->params.n_max, block_size - 1);
|
||||
this->params.n_min = std::min(this->params.n_min, block_size - 1);
|
||||
// DFlash input is [id_last, <mask> * (block_size-1)]: in-place denoising yields at most
|
||||
// block_size-1 draft tokens, DSpark yield a full block_size draft tokens
|
||||
const int32_t n_draft_max = is_dspark ? block_size : block_size - 1;
|
||||
if (this->params.n_max > n_draft_max || this->params.n_min > n_draft_max) {
|
||||
LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained block size %d -- clamping to %d\n",
|
||||
__func__, this->params.n_max, this->params.n_min, block_size, n_draft_max);
|
||||
this->params.n_max = std::min(this->params.n_max, n_draft_max);
|
||||
this->params.n_min = std::min(this->params.n_min, n_draft_max);
|
||||
}
|
||||
|
||||
batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq);
|
||||
@@ -1126,12 +1144,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
|
||||
const int32_t n = (int32_t) dp.n_past;
|
||||
|
||||
int32_t n_draft = params.n_max;
|
||||
if (dp.n_max > 0) {
|
||||
n_draft = std::min(n_draft, dp.n_max);
|
||||
}
|
||||
const int32_t n_draft = params.n_max;
|
||||
|
||||
const int32_t n_block_tokens = n_draft + 1; // id_last + n_draft * <mask>
|
||||
const int32_t n_block_tokens = n_draft + (is_dspark ? 0 : 1);
|
||||
i_block_beg[seq_id] = batch.n_tokens;
|
||||
n_block [seq_id] = n_block_tokens;
|
||||
for (int32_t i = 0; i < n_block_tokens; ++i) {
|
||||
@@ -1163,27 +1178,57 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
|
||||
|
||||
auto & result = *dp.result;
|
||||
|
||||
// greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1
|
||||
for (int32_t i = 1; i < n_block_tokens; ++i) {
|
||||
common_sampler_sample(smpl, ctx_dft, beg + i, true);
|
||||
if (is_dspark) {
|
||||
// DSpark predicts the next token from position 0 and optionally truncates
|
||||
// at the first position below the confidence threshold.
|
||||
const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr;
|
||||
|
||||
const auto * cur_p = common_sampler_get_candidates(smpl, true);
|
||||
for (int32_t i = 0; i < n_block_tokens; ++i) {
|
||||
const int32_t idx = beg + i;
|
||||
|
||||
for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
|
||||
LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
|
||||
seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p,
|
||||
common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
|
||||
if (conf && conf[(size_t) idx * n_embd_dec] < params.p_min) {
|
||||
break;
|
||||
}
|
||||
|
||||
common_sampler_sample(smpl, ctx_dft, idx, true);
|
||||
|
||||
const auto * cur_p = common_sampler_get_candidates(smpl, true);
|
||||
|
||||
for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
|
||||
LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
|
||||
seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,
|
||||
common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
|
||||
}
|
||||
|
||||
const llama_token id = cur_p->data[0].id;
|
||||
|
||||
common_sampler_accept(smpl, id, true);
|
||||
|
||||
result.push_back(id);
|
||||
}
|
||||
} else {
|
||||
// greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1
|
||||
for (int32_t i = 1; i < n_block_tokens; ++i) {
|
||||
common_sampler_sample(smpl, ctx_dft, beg + i, true);
|
||||
|
||||
const llama_token id = cur_p->data[0].id;
|
||||
const auto * cur_p = common_sampler_get_candidates(smpl, true);
|
||||
|
||||
if (cur_p->data[0].p < params.p_min) {
|
||||
break;
|
||||
for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
|
||||
LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
|
||||
seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p,
|
||||
common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
|
||||
}
|
||||
|
||||
const llama_token id = cur_p->data[0].id;
|
||||
|
||||
if (cur_p->data[0].p < params.p_min) {
|
||||
break;
|
||||
}
|
||||
|
||||
common_sampler_accept(smpl, id, true);
|
||||
|
||||
result.push_back(id);
|
||||
}
|
||||
|
||||
common_sampler_accept(smpl, id, true);
|
||||
|
||||
result.push_back(id);
|
||||
}
|
||||
|
||||
if (result.size() < (size_t) params.n_min) {
|
||||
@@ -2145,6 +2190,7 @@ std::string common_speculative_type_to_str(common_speculative_type type) {
|
||||
case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: return "draft-eagle3";
|
||||
case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: return "draft-mtp";
|
||||
case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: return "draft-dflash";
|
||||
case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: return "draft-dspark";
|
||||
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram-simple";
|
||||
case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram-map-k";
|
||||
case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v";
|
||||
@@ -2198,6 +2244,7 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) {
|
||||
case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3:
|
||||
case COMMON_SPECULATIVE_TYPE_DRAFT_MTP:
|
||||
case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH:
|
||||
case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK:
|
||||
n_max = std::max(n_max, std::max(0, spec->draft.n_max));
|
||||
break;
|
||||
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:
|
||||
@@ -2284,7 +2331,7 @@ common_speculative_init_result::common_speculative_init_result(
|
||||
std::string model_path;
|
||||
if (has_draft) {
|
||||
model_path = params.speculative.draft.mparams.path;
|
||||
LOG_TRC("%s: loading draft model '%s'\n", __func__, model_path.c_str());
|
||||
LOG_INF("%s: loading draft model '%s'\n", __func__, model_path.c_str());
|
||||
|
||||
llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams);
|
||||
if (model_dft == NULL) {
|
||||
@@ -2304,7 +2351,7 @@ common_speculative_init_result::common_speculative_init_result(
|
||||
} else if (spec_mtp) {
|
||||
model_path = params.model.path;
|
||||
|
||||
LOG_TRC("%s: creating MTP draft context against the target model '%s'\n", __func__, model_path.c_str());
|
||||
LOG_INF("%s: creating MTP draft context against the target model '%s'\n", __func__, model_path.c_str());
|
||||
|
||||
llama_context * ctx_dft = llama_init_from_model(model_tgt, cparams);
|
||||
if (ctx_dft == nullptr) {
|
||||
@@ -2342,6 +2389,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
|
||||
bool has_draft_eagle3 = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3)) && params.draft.ctx_dft != nullptr;
|
||||
bool has_draft_mtp = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_MTP)) && params.draft.ctx_dft != nullptr;
|
||||
bool has_draft_dflash = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)) && params.draft.ctx_dft != nullptr;
|
||||
bool has_draft_dspark = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)) && params.draft.ctx_dft != nullptr;
|
||||
|
||||
|
||||
|
||||
@@ -2352,7 +2400,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
|
||||
bool has_ngram_mod = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_NGRAM_MOD));
|
||||
|
||||
// when adding a new type - update here the logic above
|
||||
static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 10);
|
||||
static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 11);
|
||||
|
||||
// this list here defines the priority of the speculators
|
||||
// the one with highest priority are listed first
|
||||
@@ -2385,6 +2433,9 @@ common_speculative * common_speculative_init(common_params_speculative & params,
|
||||
if (has_draft_dflash) {
|
||||
configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, params));
|
||||
}
|
||||
if (has_draft_dspark) {
|
||||
configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, params));
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::unique_ptr<common_speculative_impl>> impls = {};
|
||||
@@ -2409,6 +2460,11 @@ common_speculative * common_speculative_init(common_params_speculative & params,
|
||||
impls.push_back(std::make_unique<common_speculative_impl_draft_dflash>(config.params, n_seq));
|
||||
break;
|
||||
}
|
||||
case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: {
|
||||
impls.push_back(std::make_unique<common_speculative_impl_draft_dflash>(
|
||||
config.params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK));
|
||||
break;
|
||||
}
|
||||
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: {
|
||||
common_ngram_map ngram_map = get_common_ngram_map(config.type, config.params.ngram_simple);
|
||||
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
#include "subproc.h"
|
||||
|
||||
bool common_subproc::is_supported() {
|
||||
#ifdef LLAMA_SUBPROCESS
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
#endif
|
||||
}
|
||||
|
||||
#ifdef LLAMA_SUBPROCESS
|
||||
|
||||
static std::vector<char *> to_cstr_vec(const std::vector<std::string> & v) {
|
||||
std::vector<char *> r;
|
||||
r.reserve(v.size() + 1);
|
||||
for (const auto & s : v) {
|
||||
r.push_back(const_cast<char *>(s.c_str()));
|
||||
}
|
||||
r.push_back(nullptr);
|
||||
return r;
|
||||
}
|
||||
|
||||
common_subproc::~common_subproc() {
|
||||
if (is_created) {
|
||||
subprocess_destroy(&proc);
|
||||
is_created = false;
|
||||
}
|
||||
}
|
||||
|
||||
bool common_subproc::create(
|
||||
const std::vector<std::string> & args,
|
||||
int options,
|
||||
const std::vector<std::string> & env,
|
||||
const char * cwd) {
|
||||
auto argv = to_cstr_vec(args);
|
||||
|
||||
int result;
|
||||
if (env.empty() && cwd == nullptr) {
|
||||
result = subprocess_create(argv.data(), options, &proc);
|
||||
} else {
|
||||
auto envp = to_cstr_vec(env);
|
||||
result = subprocess_create_ex(argv.data(), options, env.empty() ? nullptr : envp.data(), cwd, &proc);
|
||||
}
|
||||
|
||||
is_created = result == 0;
|
||||
return is_created;
|
||||
}
|
||||
|
||||
bool common_subproc::has_handle() const {
|
||||
if (!is_created) {
|
||||
return false;
|
||||
}
|
||||
#if defined(_WIN32)
|
||||
return proc.hProcess != nullptr;
|
||||
#else
|
||||
return proc.child > 0;
|
||||
#endif
|
||||
}
|
||||
|
||||
bool common_subproc::alive() {
|
||||
return is_created && subprocess_alive(&proc);
|
||||
}
|
||||
|
||||
FILE * common_subproc::stdin_file() {
|
||||
return is_created ? subprocess_stdin(&proc) : nullptr;
|
||||
}
|
||||
|
||||
FILE * common_subproc::stdout_file() {
|
||||
return is_created ? subprocess_stdout(&proc) : nullptr;
|
||||
}
|
||||
|
||||
FILE * common_subproc::stderr_file() {
|
||||
return is_created ? subprocess_stderr(&proc) : nullptr;
|
||||
}
|
||||
|
||||
void common_subproc::close_stdin() {
|
||||
if (is_created && proc.stdin_file) {
|
||||
fclose(proc.stdin_file);
|
||||
proc.stdin_file = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
void common_subproc::terminate() {
|
||||
if (has_handle()) {
|
||||
subprocess_terminate(&proc);
|
||||
}
|
||||
}
|
||||
|
||||
int common_subproc::join() {
|
||||
int exit_code = -1;
|
||||
if (is_created) {
|
||||
subprocess_join(&proc, &exit_code);
|
||||
subprocess_destroy(&proc);
|
||||
is_created = false;
|
||||
}
|
||||
return exit_code;
|
||||
}
|
||||
|
||||
#else // !LLAMA_SUBPROCESS
|
||||
|
||||
common_subproc::~common_subproc() = default;
|
||||
|
||||
bool common_subproc::create(
|
||||
const std::vector<std::string> &,
|
||||
int,
|
||||
const std::vector<std::string> &,
|
||||
const char *) {
|
||||
(void)(proc);
|
||||
(void)(is_created);
|
||||
return false;
|
||||
}
|
||||
|
||||
bool common_subproc::has_handle() const {
|
||||
return false;
|
||||
}
|
||||
|
||||
bool common_subproc::alive() {
|
||||
return false;
|
||||
}
|
||||
|
||||
FILE * common_subproc::stdin_file() {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
FILE * common_subproc::stdout_file() {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
FILE * common_subproc::stderr_file() {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void common_subproc::close_stdin() {
|
||||
}
|
||||
|
||||
void common_subproc::terminate() {
|
||||
}
|
||||
|
||||
int common_subproc::join() {
|
||||
return -1;
|
||||
}
|
||||
|
||||
#endif // LLAMA_SUBPROCESS
|
||||
@@ -0,0 +1,59 @@
|
||||
#pragma once
|
||||
|
||||
#include <atomic>
|
||||
#include <cstdio>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#ifdef LLAMA_SUBPROCESS
|
||||
#include <sheredom/subprocess.h>
|
||||
#else
|
||||
// dummy values to allow compilation when subprocess is disabled
|
||||
struct subprocess_s {};
|
||||
static constexpr int subprocess_option_no_window = 0;
|
||||
static constexpr int subprocess_option_combined_stdout_stderr = 0;
|
||||
static constexpr int subprocess_option_inherit_environment = 0;
|
||||
static constexpr int subprocess_option_search_user_path = 0;
|
||||
#endif
|
||||
|
||||
// RAII-style wrapper around https://github.com/sheredom/subprocess.h,
|
||||
// exposing method calls instead of free functions operating on subprocess_s.
|
||||
struct common_subproc {
|
||||
common_subproc() = default;
|
||||
~common_subproc();
|
||||
|
||||
common_subproc(const common_subproc &) = delete;
|
||||
common_subproc & operator=(const common_subproc &) = delete;
|
||||
|
||||
// spawn a child process; if env is non-empty it replaces the child's environment
|
||||
// (do not combine with subprocess_option_inherit_environment)
|
||||
bool create(
|
||||
const std::vector<std::string> & args,
|
||||
int options,
|
||||
const std::vector<std::string> & env = {},
|
||||
const char * cwd = nullptr);
|
||||
|
||||
bool alive();
|
||||
|
||||
// true if LLAMA_SUBPROCESS was enabled at build time; when false, create() always fails
|
||||
static bool is_supported();
|
||||
|
||||
FILE * stdin_file();
|
||||
FILE * stdout_file();
|
||||
FILE * stderr_file();
|
||||
|
||||
// close stdin and detach it from the process, so a later join()/destroy() won't double-close it;
|
||||
// use this after writing all input to signal EOF to the child while it's still running
|
||||
void close_stdin();
|
||||
|
||||
void terminate();
|
||||
|
||||
// wait for the process to exit, release the underlying handle and return its exit code
|
||||
int join();
|
||||
|
||||
private:
|
||||
subprocess_s proc {};
|
||||
std::atomic<bool> is_created{false};
|
||||
|
||||
bool has_handle() const;
|
||||
};
|
||||
+123
@@ -0,0 +1,123 @@
|
||||
#include "trie.h"
|
||||
|
||||
#include "unicode.h"
|
||||
|
||||
#include <deque>
|
||||
|
||||
common_trie::match_result common_trie::check_at(std::string_view sv, size_t start_pos) const {
|
||||
size_t current = 0; // Start at root
|
||||
size_t pos = start_pos;
|
||||
|
||||
// LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str());
|
||||
|
||||
while (pos < sv.size()) {
|
||||
auto result = common_parse_utf8_codepoint(sv, pos);
|
||||
if (result.status != utf8_parse_result::SUCCESS) {
|
||||
break;
|
||||
}
|
||||
|
||||
auto it = nodes[current].children.find(result.codepoint);
|
||||
if (it == nodes[current].children.end()) {
|
||||
// Can't continue matching
|
||||
return match_result{match_result::NO_MATCH};
|
||||
}
|
||||
|
||||
current = it->second;
|
||||
pos += result.bytes_consumed;
|
||||
|
||||
// Check if we've matched a complete word
|
||||
if (nodes[current].pattern >= 0) {
|
||||
return match_result{match_result::COMPLETE_MATCH};
|
||||
}
|
||||
}
|
||||
|
||||
// Reached end of input while still in the trie (not at root)
|
||||
if (current != 0) {
|
||||
// We're in the middle of a potential match
|
||||
return match_result{match_result::PARTIAL_MATCH};
|
||||
}
|
||||
|
||||
// Reached end at root (no match)
|
||||
return match_result{match_result::NO_MATCH};
|
||||
}
|
||||
|
||||
int32_t common_trie::insert(const std::string & word) {
|
||||
std::vector<uint32_t> symbols;
|
||||
size_t pos = 0;
|
||||
while (pos < word.length()) {
|
||||
auto result = common_parse_utf8_codepoint(word, pos);
|
||||
if (result.status != utf8_parse_result::SUCCESS) {
|
||||
break;
|
||||
}
|
||||
|
||||
symbols.push_back(result.codepoint);
|
||||
pos += result.bytes_consumed;
|
||||
}
|
||||
return insert(symbols);
|
||||
}
|
||||
|
||||
int32_t common_trie::insert(const std::vector<uint32_t> & symbols) {
|
||||
size_t current = 0;
|
||||
for (uint32_t ch : symbols) {
|
||||
auto it = nodes[current].children.find(ch);
|
||||
if (it == nodes[current].children.end()) {
|
||||
size_t child = create_node();
|
||||
nodes[current].children[ch] = child;
|
||||
current = child;
|
||||
} else {
|
||||
current = it->second;
|
||||
}
|
||||
}
|
||||
if (nodes[current].pattern < 0) {
|
||||
nodes[current].pattern = n_patterns++;
|
||||
}
|
||||
return nodes[current].pattern;
|
||||
}
|
||||
|
||||
common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) {
|
||||
const auto & nodes = t.nodes;
|
||||
const size_t n = nodes.size();
|
||||
|
||||
fail.assign(n, 0);
|
||||
order.reserve(n);
|
||||
|
||||
std::deque<size_t> queue{ 0 };
|
||||
while (!queue.empty()) {
|
||||
size_t u = queue.front();
|
||||
queue.pop_front();
|
||||
order.push_back(u);
|
||||
for (const auto & [ch, v] : nodes[u].children) {
|
||||
if (u != 0) {
|
||||
size_t f = fail[u];
|
||||
while (f && nodes[f].children.find(ch) == nodes[f].children.end()) {
|
||||
f = fail[f];
|
||||
}
|
||||
auto it = nodes[f].children.find(ch);
|
||||
fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0;
|
||||
}
|
||||
queue.push_back(v);
|
||||
}
|
||||
}
|
||||
|
||||
// fail[u] points to a strictly shorter suffix, so the first pattern found on
|
||||
// the fail chain (including u itself) is the longest pattern ending at u
|
||||
match.assign(n, -1);
|
||||
for (size_t u : order) {
|
||||
match[u] = nodes[u].pattern >= 0 ? nodes[u].pattern : (u != 0 ? match[fail[u]] : -1);
|
||||
}
|
||||
|
||||
for (const auto & node : nodes) {
|
||||
for (const auto & [ch, v] : node.children) {
|
||||
alphabet.insert(ch);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
size_t common_aho_corasick::next(size_t state, uint32_t ch) const {
|
||||
const auto & nodes = t.nodes;
|
||||
while (state && nodes[state].children.find(ch) == nodes[state].children.end()) {
|
||||
state = fail[state];
|
||||
}
|
||||
auto it = nodes[state].children.find(ch);
|
||||
return it != nodes[state].children.end() ? it->second : 0;
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
#pragma once
|
||||
|
||||
#include <cstdint>
|
||||
#include <map>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
// Trie for matching multiple literals.
|
||||
// This is used in common_peg_until_parser and to build a GBNF exclusion grammar
|
||||
struct common_trie {
|
||||
struct node {
|
||||
std::map<uint32_t, size_t> children; // Use uint32_t to store Unicode codepoints
|
||||
int32_t pattern = -1; // index of the pattern ending at this node, -1 if none
|
||||
};
|
||||
|
||||
std::vector<node> nodes;
|
||||
|
||||
common_trie() {
|
||||
create_node(); // root node
|
||||
}
|
||||
|
||||
common_trie(const std::vector<std::string> & words) : common_trie() {
|
||||
for (const auto & w : words) {
|
||||
insert(w);
|
||||
}
|
||||
}
|
||||
|
||||
enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH };
|
||||
|
||||
// Check if a delimiter starts at the given position
|
||||
match_result check_at(std::string_view sv, size_t start_pos) const;
|
||||
|
||||
// Insert a word as a sequence of Unicode codepoints, returns its pattern index
|
||||
int32_t insert(const std::string & word);
|
||||
|
||||
// Insert a raw symbol sequence, returns its pattern index (insertion order,
|
||||
// duplicates keep the first index)
|
||||
int32_t insert(const std::vector<uint32_t> & symbols);
|
||||
|
||||
private:
|
||||
int32_t n_patterns = 0;
|
||||
|
||||
size_t create_node() {
|
||||
size_t index = nodes.size();
|
||||
nodes.emplace_back();
|
||||
return index;
|
||||
}
|
||||
};
|
||||
|
||||
// Aho-Corasick automaton
|
||||
struct common_aho_corasick {
|
||||
common_trie t;
|
||||
std::vector<size_t> fail; // failure links
|
||||
std::vector<size_t> order; // states in BFS order
|
||||
std::vector<int32_t> match; // longest pattern ending at each state (directly or via a suffix link), -1 if none
|
||||
std::set<uint32_t> alphabet; // every character with a transition
|
||||
|
||||
common_aho_corasick(common_trie trie);
|
||||
|
||||
common_aho_corasick(const std::vector<std::string> & strings)
|
||||
: common_aho_corasick(common_trie(strings)) {}
|
||||
|
||||
size_t num_states() const { return t.nodes.size(); }
|
||||
bool is_terminal(size_t s) const { return match[s] >= 0; }
|
||||
|
||||
// index of the longest pattern ending at this state, -1 if none
|
||||
int32_t match_pattern(size_t s) const { return match[s]; }
|
||||
|
||||
// follow failure links until a transition on `ch` exists.
|
||||
size_t next(size_t state, uint32_t ch) const;
|
||||
};
|
||||
@@ -53,6 +53,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"DeepseekV3ForCausalLM": "deepseek",
|
||||
"DeepseekV32ForCausalLM": "deepseek",
|
||||
"DFlashDraftModel": "qwen",
|
||||
"Qwen3DSparkModel": "qwen",
|
||||
"DeepseekV4ForCausalLM": "deepseek",
|
||||
"DistilBertForMaskedLM": "bert",
|
||||
"DistilBertForSequenceClassification": "bert",
|
||||
@@ -158,6 +159,8 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"MiniCPMForCausalLM": "minicpm",
|
||||
"MiniCPMV4_6ForConditionalGeneration": "minicpm",
|
||||
"MiniMaxM2ForCausalLM": "minimax",
|
||||
"MiniMaxM3SparseForCausalLM": "minimax",
|
||||
"MiniMaxM3SparseForConditionalGeneration": "minimax",
|
||||
"Ministral3ForCausalLM": "mistral3",
|
||||
"Mistral3ForConditionalGeneration": "mistral3",
|
||||
"MistralForCausalLM": "llama",
|
||||
@@ -165,6 +168,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"ModernBertForMaskedLM": "bert",
|
||||
"ModernBertForSequenceClassification": "bert",
|
||||
"ModernBertModel": "bert",
|
||||
"NanbeigeForCausalLM": "nanbeige",
|
||||
"NemotronForCausalLM": "nemotron",
|
||||
"NemotronHForCausalLM": "nemotron",
|
||||
"NeoBERT": "bert",
|
||||
@@ -267,6 +271,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
|
||||
"Gemma4UnifiedForConditionalGeneration": "gemma",
|
||||
"Glm4vForConditionalGeneration": "qwen3vl",
|
||||
"Glm4vMoeForConditionalGeneration": "qwen3vl",
|
||||
"Glm5vForConditionalGeneration": "kimivl",
|
||||
"GlmOcrForConditionalGeneration": "qwen3vl",
|
||||
"GlmasrModel": "ultravox",
|
||||
"Granite4VisionForConditionalGeneration": "granite",
|
||||
@@ -285,6 +290,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
|
||||
"LlavaForConditionalGeneration": "llava",
|
||||
"MERaLiON2ForConditionalGeneration": "ultravox",
|
||||
"MiMoV2ForCausalLM": "mimo",
|
||||
"MiniMaxM3SparseForConditionalGeneration": "minimax",
|
||||
"MiniCPMV4_6ForConditionalGeneration": "minicpm",
|
||||
"Mistral3ForConditionalGeneration": "llava",
|
||||
"NemotronH_Nano_VL_V2": "nemotron",
|
||||
|
||||
+2
-2
@@ -1156,7 +1156,7 @@ class TextModel(ModelBase):
|
||||
or "projector." in name or "pre_mm_projector_norm" in name \
|
||||
or "image_newline" in name or "view_seperator" in name \
|
||||
or "patch_embed" in name or "patch_embedding" in name \
|
||||
or "patch_merger." in name or "model.connector." in name:
|
||||
or "patch_merger." in name or "patch_merge_mlp." in name or "model.connector." in name:
|
||||
return None
|
||||
|
||||
return super().filter_tensors(item)
|
||||
@@ -1203,7 +1203,7 @@ class TextModel(ModelBase):
|
||||
self.gguf_writer.add_embedding_length(n_embd)
|
||||
logger.info(f"gguf: embedding length = {n_embd}")
|
||||
|
||||
if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None:
|
||||
if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None:
|
||||
self.gguf_writer.add_feed_forward_length(n_ff)
|
||||
logger.info(f"gguf: feed forward length = {n_ff}")
|
||||
|
||||
|
||||
@@ -237,6 +237,9 @@ class GlmMoeDsaModel(DeepseekV2Model):
|
||||
self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"])
|
||||
self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"])
|
||||
self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"])
|
||||
if (indexer_types := self.hparams.get("indexer_types")) is not None:
|
||||
indexer_types = [t == "full" for t in indexer_types]
|
||||
self.gguf_writer.add_indexer_types(indexer_types)
|
||||
|
||||
|
||||
@ModelBase.register("SolarOpenForCausalLM")
|
||||
|
||||
@@ -152,3 +152,19 @@ class KimiK25Model(MmprojModel):
|
||||
name = name.replace(".proj.2.", ".proj.linear_2.")
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
|
||||
@ModelBase.register("Glm5vForConditionalGeneration")
|
||||
class Glm5vModel(KimiK25Model):
|
||||
"""GLM-5.2-Vision MoonViT3d encoder and projector
|
||||
|
||||
Uses the same vision encoder and projector as Kimi-K2.5, so it reuses the
|
||||
kimik25 projector type. The image begin/end tokens differ, but they are
|
||||
resolved at runtime from the text model vocab.
|
||||
"""
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
if name.startswith("mm_projector.linear_"):
|
||||
name = name.replace("mm_projector.linear_", "mm_projector.proj.linear_", 1)
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
+16
-2
@@ -69,9 +69,14 @@ class LlamaModel(TextModel):
|
||||
target_config = {**target_config, **target_config["text_config"]}
|
||||
self.target_vocab_size = target_config["vocab_size"]
|
||||
|
||||
# target_layers: derived from target model layer count (low/mid/high)
|
||||
# target_layers: use the eagle3 config's explicit aux hidden-state layer ids
|
||||
# if present, else derive from the target layer count.
|
||||
target_num_layers = target_config["num_hidden_layers"]
|
||||
target_layers = [2, target_num_layers // 2, target_num_layers - 3]
|
||||
aux_layer_ids = eagle3_raw_config.get("eagle_aux_hidden_state_layer_ids")
|
||||
if aux_layer_ids:
|
||||
target_layers = aux_layer_ids
|
||||
else:
|
||||
target_layers = [2, target_num_layers // 2, target_num_layers - 3]
|
||||
logger.info(f"EAGLE-3: target_layers = {target_layers} (target model has {target_num_layers} layers)")
|
||||
self.gguf_writer.add_target_layers(target_layers)
|
||||
|
||||
@@ -90,6 +95,12 @@ class LlamaModel(TextModel):
|
||||
logger.info(f"EAGLE-3: norm_before_residual = {norm_before_residual}")
|
||||
self.gguf_writer.add_norm_before_residual(norm_before_residual)
|
||||
|
||||
# norm_before_fc: RMSNorm applied to the fused target features before the
|
||||
# fc projection (e.g. nvidia/gpt-oss-120b-Eagle3-v3)
|
||||
norm_before_fc = eagle3_raw_config.get("norm_before_fc", False)
|
||||
logger.info(f"EAGLE-3: norm_before_fc = {norm_before_fc}")
|
||||
self.gguf_writer.add_norm_before_fc(norm_before_fc)
|
||||
|
||||
def set_vocab(self):
|
||||
# eagle3: use tokenizer from target model if provided
|
||||
original_dir_model = None
|
||||
@@ -222,6 +233,9 @@ class LlamaModel(TextModel):
|
||||
if name == "fc.weight":
|
||||
yield (name, data_torch)
|
||||
return
|
||||
if name == "input_norm.weight":
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ENC_OUTPUT_NORM), data_torch)
|
||||
return
|
||||
if name == "d2t":
|
||||
# store for manual int64 handling in prepare_tensors (avoid F32 conversion)
|
||||
if not hasattr(self, '_eagle3_int_tensors'):
|
||||
|
||||
+114
-9
@@ -1,8 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
|
||||
from typing import Callable, TYPE_CHECKING
|
||||
from typing import Any, Callable, Iterable, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
@@ -229,7 +230,13 @@ class MimoV2Model(TextModel):
|
||||
|
||||
|
||||
@ModelBase.register("MiMoV2ForCausalLM")
|
||||
class MiMoV2VisionModel(MmprojModel):
|
||||
class MiMoV2VisionAudioModel(MmprojModel):
|
||||
has_audio_encoder = True
|
||||
|
||||
_audio_tok_hparams: dict[str, Any] | None = None
|
||||
_rvq_codebook_sizes: list[int] | None = None
|
||||
_code_embd: dict[int, Tensor] | None = None
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
assert self.hparams_vision is not None
|
||||
@@ -253,10 +260,22 @@ class MiMoV2VisionModel(MmprojModel):
|
||||
self.visual_token_window_size = int(hp.get("visual_token_window_size", -1))
|
||||
self.use_sink = bool(hp.get("use_sink", False))
|
||||
|
||||
def get_audio_config(self) -> dict[str, Any] | None:
|
||||
if self._audio_tok_hparams is None:
|
||||
path = self.dir_model / "audio_tokenizer" / "config.json"
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
# aliases so MmprojModel.find_aparam() / n_block_keys can resolve them
|
||||
cfg["hidden_size"] = cfg["d_model"]
|
||||
cfg["intermediate_size"] = cfg["encoder_ffn_dim"]
|
||||
cfg["num_attention_heads"] = cfg["encoder_attention_heads"]
|
||||
self._audio_tok_hparams = cfg
|
||||
return self._audio_tok_hparams
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
|
||||
self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MIMOVL)
|
||||
self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.MIMOVL)
|
||||
self.gguf_writer.add_vision_use_silu(True)
|
||||
self.gguf_writer.add_vision_head_count_kv(self.num_kv_heads)
|
||||
self.gguf_writer.add_vision_spatial_merge_size(self.spatial_merge_size)
|
||||
@@ -266,19 +285,45 @@ class MiMoV2VisionModel(MmprojModel):
|
||||
self.gguf_writer.add_vision_min_pixels(int(self.preprocessor_config["min_pixels"]))
|
||||
self.gguf_writer.add_vision_max_pixels(int(self.preprocessor_config["max_pixels"]))
|
||||
|
||||
assert self.hparams_audio is not None
|
||||
self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.MIMO_AUDIO)
|
||||
self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["n_mels"])
|
||||
self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams_audio.get("layer_norm_eps", 1e-5))
|
||||
|
||||
assert self._rvq_codebook_sizes is not None
|
||||
self.gguf_writer.add_audio_rvq_num_quantizers(len(self._rvq_codebook_sizes))
|
||||
self.gguf_writer.add_audio_rvq_codebook_size(self._rvq_codebook_sizes)
|
||||
|
||||
n_layer = self.hparams_audio["encoder_layers"]
|
||||
swa_per_block = self.hparams_audio.get("swa_per_block", 1)
|
||||
if self.hparams_audio.get("hybrid_attention") and swa_per_block > 1:
|
||||
wa_pattern = [0 if i % swa_per_block < swa_per_block - 1 else -1 for i in range(n_layer)]
|
||||
else:
|
||||
wa_pattern = [-1] * n_layer
|
||||
self.gguf_writer.add_audio_wa_pattern_mode(wa_pattern)
|
||||
self.gguf_writer.add_audio_window_size(int(self.hparams_audio["encoder_attn_window_size"][0]))
|
||||
|
||||
audio_cfg = self.global_config["audio_config"]
|
||||
self.gguf_writer.add_audio_local_block_count(int(audio_cfg["input_local_layers"]))
|
||||
self.gguf_writer.add_audio_local_group_size(int(audio_cfg["group_size"]))
|
||||
|
||||
def tensor_force_quant(self, name, new_name, bid, n_dims):
|
||||
# Sinks must be F32: any sink-style softmax/mask add in ggml requires
|
||||
# F32, and we fold sinks into a host-built F32 mask at encode time.
|
||||
if new_name.endswith(".attn_sinks"):
|
||||
# for audio encoder: keep codebook in F32
|
||||
if new_name in (
|
||||
gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_ENC_RVQ_CODEBOOK] + ".weight",
|
||||
gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_MM_CODE_EMBD] + ".weight",
|
||||
):
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
if ("encoder.conv" in name or "encoder.down_sample_layer" in name) and name.endswith(".weight"):
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
return super().tensor_force_quant(name, new_name, bid, n_dims)
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
name, _ = item
|
||||
if not name.startswith("visual."):
|
||||
return None
|
||||
return super().filter_tensors(item)
|
||||
if name.startswith("visual.") or name.startswith("speech_embeddings.") or name.startswith("audio_encoder."):
|
||||
return super().filter_tensors(item)
|
||||
return None
|
||||
|
||||
def modify_tensors(self, data_torch, name, bid):
|
||||
# Conv3D patch embed: split along the temporal axis (kt=2) into two Conv2D
|
||||
@@ -292,4 +337,64 @@ class MiMoV2VisionModel(MmprojModel):
|
||||
yield (embd_name + ".weight.1", data_torch[:, :, 1, ...])
|
||||
return
|
||||
|
||||
if m := re.match(r"^speech_embeddings\.(\d+)\.weight$", name):
|
||||
if self._code_embd is None:
|
||||
self._code_embd = {}
|
||||
self._code_embd[int(m.group(1))] = data_torch
|
||||
|
||||
n_channels = int(self.global_config["audio_config"]["audio_channels"])
|
||||
if len(self._code_embd) < n_channels:
|
||||
return
|
||||
merged = torch.stack([self._code_embd.pop(i) for i in range(n_channels)], dim=0)
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_MM_CODE_EMBD), merged)
|
||||
return
|
||||
|
||||
if "conv1.bias" in name or "conv2.bias" in name:
|
||||
# transpose conv1/conv2 bias so it broadcasts against [n_frames, C_out, 1]
|
||||
data_torch = data_torch.unsqueeze(-1)
|
||||
|
||||
if name == "audio_encoder.projection.mlp.0.weight":
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_MMPROJ, 1), data_torch)
|
||||
return
|
||||
if name == "audio_encoder.projection.mlp.2.weight":
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_MMPROJ, 2), data_torch)
|
||||
return
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
|
||||
# note: audio encoder is in its own subdir "audio_tokenizer"
|
||||
from safetensors.torch import load_file
|
||||
|
||||
tok_dir = self.dir_model / "audio_tokenizer"
|
||||
state_dict = load_file(tok_dir / "model.safetensors")
|
||||
|
||||
codebook_re = re.compile(r"^encoder\.quantizer\.vq\.layers\.(\d+)\._codebook\.embed$")
|
||||
codebooks: dict[int, Tensor] = {}
|
||||
|
||||
# EMA/training-only RVQ buffers - not needed for inference (nearest-codebook
|
||||
# lookup only reads "_codebook.embed")
|
||||
skip_suffixes = (
|
||||
"_codebook.cluster_size",
|
||||
"_codebook.embed_avg",
|
||||
"_codebook.inited",
|
||||
)
|
||||
for name, tensor in state_dict.items():
|
||||
if name.endswith(skip_suffixes):
|
||||
continue
|
||||
if m := codebook_re.match(name):
|
||||
codebooks[int(m.group(1))] = tensor
|
||||
continue
|
||||
yield name, tensor
|
||||
|
||||
# gather codebooks and merge into 3D tensor, similar to MoE MLP tensors
|
||||
n_q = len(codebooks)
|
||||
ordered = [codebooks[i] for i in range(n_q)]
|
||||
self._rvq_codebook_sizes = [int(cb.shape[0]) for cb in ordered]
|
||||
max_bins = max(self._rvq_codebook_sizes)
|
||||
dim = ordered[0].shape[1]
|
||||
merged = ordered[0].new_zeros(n_q, max_bins, dim)
|
||||
for i, cb in enumerate(ordered):
|
||||
merged[i, : cb.shape[0], :] = cb
|
||||
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_ENC_RVQ_CODEBOOK), merged)
|
||||
|
||||
+117
-2
@@ -7,7 +7,7 @@ import torch
|
||||
if TYPE_CHECKING:
|
||||
from torch import Tensor
|
||||
|
||||
from .base import ModelBase, TextModel, gguf
|
||||
from .base import ModelBase, TextModel, MmprojModel, gguf
|
||||
|
||||
|
||||
@ModelBase.register("MiniMaxM2ForCausalLM")
|
||||
@@ -23,7 +23,7 @@ class MiniMaxM2Model(TextModel):
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
|
||||
# merge expert weights
|
||||
if 'experts' in name:
|
||||
if "block_sparse_moe.experts." in name:
|
||||
n_experts = self.find_hparam(["num_local_experts", "num_experts"])
|
||||
assert bid is not None
|
||||
|
||||
@@ -52,3 +52,118 @@ class MiniMaxM2Model(TextModel):
|
||||
return
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
|
||||
@ModelBase.register("MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration")
|
||||
class MiniMaxM3Model(MiniMaxM2Model):
|
||||
model_arch = gguf.MODEL_ARCH.MINIMAXM3
|
||||
|
||||
def tensor_force_quant(self, name, new_name, bid, n_dims):
|
||||
if ".indexer." in new_name:
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
return super().tensor_force_quant(name, new_name, bid, n_dims)
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
|
||||
self.gguf_writer.add_expert_shared_count(self.find_hparam(["n_shared_experts"]))
|
||||
self.gguf_writer.add_expert_weights_scale(self.find_hparam(["routed_scaling_factor"]))
|
||||
self.gguf_writer.add_expert_weights_norm(True)
|
||||
|
||||
sac = self.find_hparam(["sparse_attention_config"])
|
||||
self.gguf_writer.add_indexer_head_count(sac["sparse_num_index_heads"])
|
||||
self.gguf_writer.add_indexer_key_length(sac["sparse_index_dim"])
|
||||
self.gguf_writer.add_indexer_top_k(sac["sparse_topk_blocks"])
|
||||
self.gguf_writer.add_indexer_block_size(sac["sparse_block_size"])
|
||||
self.gguf_writer.add_indexer_local_blocks(sac["sparse_local_block"])
|
||||
|
||||
moe_layer_freq = self.find_hparam(["moe_layer_freq"])
|
||||
n_dense = 0
|
||||
for v in moe_layer_freq:
|
||||
if v == 0:
|
||||
n_dense += 1
|
||||
else:
|
||||
break
|
||||
self.gguf_writer.add_leading_dense_block_count(n_dense)
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
|
||||
# Gemma-style (1 + w) RMSNorm: bake the +1 in so llama.cpp can use plain RMSNorm
|
||||
if name.endswith("norm.weight"):
|
||||
data_torch = data_torch + 1.0
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
|
||||
@ModelBase.register("MiniMaxM3SparseForConditionalGeneration", "MiniMaxM3VLForConditionalGeneration")
|
||||
class MiniMaxM3VisionModel(MmprojModel):
|
||||
@classmethod
|
||||
def filter_tensors(cls, item):
|
||||
name, gen = item
|
||||
# keep only the vision-side tensors; text / mtp / sparse-index are dropped
|
||||
if not name.startswith(("vision_tower.", "multi_modal_projector.", "patch_merge_mlp.")):
|
||||
return None
|
||||
return super().filter_tensors((name, gen))
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
assert self.hparams_vision is not None
|
||||
|
||||
self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MINIMAXM3)
|
||||
self.gguf_writer.add_vision_use_gelu(True)
|
||||
|
||||
# the ViT carries its own LayerNorm eps (text tower uses a different one)
|
||||
self.gguf_writer.add_vision_attention_layernorm_eps(
|
||||
self.hparams_vision.get("layer_norm_eps", 1e-5)
|
||||
)
|
||||
|
||||
comp = self.hparams_vision.get("img_token_compression_config", {})
|
||||
merge_size = comp.get("spatial_merge_size", 2)
|
||||
self.gguf_writer.add_vision_spatial_merge_size(int(merge_size))
|
||||
|
||||
def modify_tensors(self, data_torch, name, bid):
|
||||
assert self.hparams_vision is not None
|
||||
|
||||
# Conv3d patch embed -> Conv2d slices
|
||||
if name == "vision_tower.vision_model.embeddings.patch_embedding.weight":
|
||||
if data_torch.ndim != 5:
|
||||
raise ValueError(f"unexpected patch_embedding rank {data_torch.ndim} for {name}")
|
||||
kt = data_torch.shape[2]
|
||||
base = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_ENC_EMBD_PATCH]
|
||||
for t in range(kt):
|
||||
suffix = ".weight" if t == 0 else f".weight.{t}"
|
||||
yield (base + suffix, data_torch[:, :, t, ...])
|
||||
return
|
||||
|
||||
# Permute ViT q/k. HF [Ta Ha Wa | Tb Hb Wb | pad] reorder to [Ta Tb | Ha Hb | Wa Wb | pad].
|
||||
for new_name, tensor in super().modify_tensors(data_torch, name, bid):
|
||||
if ".attn_q." in new_name or ".attn_k." in new_name:
|
||||
tensor = self._permute_vit_qk(tensor, new_name)
|
||||
yield new_name, tensor
|
||||
|
||||
def _permute_vit_qk(self, t: "Tensor", new_name: str) -> "Tensor":
|
||||
assert self.hparams_vision is not None
|
||||
n_head = self.hparams_vision["num_attention_heads"]
|
||||
d_head = t.shape[0] // n_head
|
||||
axis_dim = 2 * ((2 * (d_head // 2) // 3) // 2)
|
||||
ah = axis_dim // 2
|
||||
half = 3 * ah
|
||||
perm = []
|
||||
perm += list(range(0, ah))
|
||||
perm += list(range(half, half + ah))
|
||||
perm += list(range(ah, 2 * ah))
|
||||
perm += list(range(half + ah, half + 2 * ah))
|
||||
perm += list(range(2 * ah, 3 * ah))
|
||||
perm += list(range(half + 2 * ah, half + 3 * ah))
|
||||
perm += list(range(2 * half, d_head))
|
||||
|
||||
assert axis_dim % 2 == 0
|
||||
assert 3 * axis_dim <= d_head
|
||||
assert len(perm) == d_head
|
||||
assert sorted(perm) == list(range(d_head)), "perm is not a bijection of d_head"
|
||||
assert t.shape[0] == n_head * d_head, f"{new_name}: {t.shape[0]} != {n_head}*{d_head}"
|
||||
assert d_head == 80
|
||||
|
||||
idx = torch.tensor(perm, dtype=torch.long)
|
||||
if t.ndim == 2:
|
||||
return t.reshape(n_head, d_head, t.shape[1])[:, idx, :].reshape(t.shape)
|
||||
return t.reshape(n_head, d_head)[:, idx].reshape(t.shape)
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from .base import ModelBase, gguf, logger
|
||||
from .llama import LlamaModel
|
||||
|
||||
|
||||
@ModelBase.register("NanbeigeForCausalLM")
|
||||
class NanbeigeModel(LlamaModel):
|
||||
model_arch = gguf.MODEL_ARCH.NANBEIGE
|
||||
undo_permute = True
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
hparams = self.hparams
|
||||
|
||||
n_loops = int(hparams.get("num_loops", 1) or 1)
|
||||
if n_loops < 1:
|
||||
n_loops = 1
|
||||
self.gguf_writer.add_num_loops(n_loops)
|
||||
logger.info(f"gguf: num_loops = {n_loops}")
|
||||
|
||||
skip_loop_final_norm = bool(hparams.get("skip_loop_final_norm", False))
|
||||
self.gguf_writer.add_skip_loop_final_norm(skip_loop_final_norm)
|
||||
logger.info(f"gguf: skip_loop_final_norm = {skip_loop_final_norm}")
|
||||
+50
-7
@@ -39,28 +39,48 @@ class NemotronNanoV2VLModel(MmprojModel):
|
||||
}
|
||||
return vision_config
|
||||
|
||||
def get_audio_config(self) -> dict[str, Any] | None:
|
||||
return self.global_config.get("sound_config")
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
if "image_mean" not in self.preprocessor_config:
|
||||
self.preprocessor_config["image_mean"] = [0.485, 0.456, 0.406]
|
||||
if "image_std" not in self.preprocessor_config:
|
||||
self.preprocessor_config["image_std"] = [0.229, 0.224, 0.225]
|
||||
|
||||
if self.hparams_audio is not None:
|
||||
self.has_vision_encoder = True
|
||||
self.has_audio_encoder = True
|
||||
self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"])
|
||||
self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
|
||||
self.gguf_writer.add_audio_subsampling_factor(self.hparams_audio["subsampling_factor"])
|
||||
self.gguf_writer.add_audio_conv_kernel_size(self.hparams_audio["conv_kernel_size"])
|
||||
self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.PARAKEET)
|
||||
self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)
|
||||
else:
|
||||
self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)
|
||||
|
||||
super().set_gguf_parameters()
|
||||
hparams = self.global_config
|
||||
self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)
|
||||
self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)
|
||||
self.gguf_writer.add_vision_use_gelu(True)
|
||||
downsample_ratio = hparams.get("downsample_ratio", 0.5)
|
||||
self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / downsample_ratio))
|
||||
|
||||
def tensor_force_quant(self, name, new_name, bid, n_dims):
|
||||
if ".position_embd." in new_name or "pos_embed" in new_name:
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
if "sound_encoder" in name or new_name.startswith("mm.a."):
|
||||
if "bias" in new_name or "norm" in new_name:
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
if "conv" in new_name and "weight" in new_name:
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
|
||||
return super().tensor_force_quant(name, new_name, bid, n_dims)
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
name, gen = item
|
||||
if (titem := super().filter_tensors(item)) is None:
|
||||
return None
|
||||
name, gen = titem
|
||||
|
||||
if "input_conditioner" in name:
|
||||
return None
|
||||
@@ -69,14 +89,18 @@ class NemotronNanoV2VLModel(MmprojModel):
|
||||
if "radio_model.model.patch_generator.video_embedder" in name:
|
||||
return None
|
||||
|
||||
if not name.startswith("vision_model.radio_model.model.") and not name.startswith("mlp1."):
|
||||
if not name.startswith(("vision_model.radio_model.model.", "mlp1.", "sound_encoder.", "sound_projection.")):
|
||||
return None
|
||||
|
||||
if "patch_generator.pos_embed" in name:
|
||||
if not name.endswith(".weight"):
|
||||
name += ".weight"
|
||||
|
||||
return super().filter_tensors((name, gen))
|
||||
# num_batches is only used for training not inference.
|
||||
if "conv.norm" in name and "num_batches" in name:
|
||||
return None
|
||||
|
||||
return name, gen
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
# RADIO's pos_embed doesn't have .weight suffix, but clip.cpp expects it
|
||||
@@ -104,7 +128,26 @@ class NemotronNanoV2VLModel(MmprojModel):
|
||||
n_embd = self.hparams["hidden_size"]
|
||||
data_torch = data_torch.reshape(n_embd, 3, patch_size, patch_size)
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
if "depthwise_conv.weight" in name:
|
||||
data_torch = data_torch.unsqueeze(-1)
|
||||
data_torch = data_torch.permute(3, 1, 0, 2).contiguous()
|
||||
|
||||
if "pointwise_conv" in name and name.endswith(".weight"):
|
||||
if len(data_torch.shape) == 3 and data_torch.shape[2] == 1:
|
||||
data_torch = data_torch.reshape(data_torch.shape[0], data_torch.shape[1])
|
||||
|
||||
if "subsampling.layers" in name and name.endswith(".bias"):
|
||||
if len(data_torch.shape) == 1:
|
||||
data_torch = data_torch.reshape(1, -1, 1, 1)
|
||||
|
||||
if "pointwise_conv" in name and name.endswith(".bias"):
|
||||
if len(data_torch.shape) == 1:
|
||||
data_torch = data_torch.reshape(1, -1, 1, 1)
|
||||
|
||||
for mapped_name, tensor in super().modify_tensors(data_torch, name, bid):
|
||||
if name.startswith("sound_projection.") and mapped_name.startswith("mm.model.mlp."):
|
||||
mapped_name = mapped_name.replace("mm.model.mlp.", "mm.a.mlp.")
|
||||
yield mapped_name, tensor
|
||||
|
||||
|
||||
@ModelBase.register("NemotronForCausalLM")
|
||||
|
||||
@@ -688,3 +688,23 @@ class DFlashModel(Qwen3Model):
|
||||
if not name.startswith("model."):
|
||||
name = "model." + name
|
||||
return super().filter_tensors((name, gen))
|
||||
|
||||
|
||||
@ModelBase.register("Qwen3DSparkModel")
|
||||
class DSparkModel(DFlashModel):
|
||||
# DSpark = DFlash + a semi-autoregressive Markov head
|
||||
model_arch = gguf.MODEL_ARCH.DFLASH
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
# normalize the flat DeepSpec schema to DFlash's nested dflash_config
|
||||
self.hparams.setdefault("dflash_config", {
|
||||
k: self.hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in self.hparams
|
||||
})
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
name, gen = item
|
||||
if name.endswith(("embed_tokens.weight", "lm_head.weight")):
|
||||
return None
|
||||
return super().filter_tensors((name, gen))
|
||||
|
||||
@@ -98,6 +98,24 @@ The OpenCL backend has the following CMake options that control the behavior of
|
||||
| `GGML_OPENCL_USE_ADRENO_KERNELS` | `ON` | Use kernels optimized for Adreno. |
|
||||
| `GGML_OPENCL_USE_ADRENO_BIN_KERNELS` | `OFF` | Allow using binary kernel lib for Adreno. |
|
||||
|
||||
## Program Binary Cache
|
||||
|
||||
Compiled `cl_program` binaries are cached on disk, so subsequent runs skip the expensive
|
||||
compile-from-source step when nothing relevant has changed (kernel source, compile options,
|
||||
device, driver, or platform version).
|
||||
|
||||
The cache is controlled with the `GGML_OPENCL_KERNEL_CACHE_DIR` environment variable:
|
||||
|
||||
| Value | Behavior |
|
||||
|:---------------------------------------|:-----------------------------------------------|
|
||||
| unset / empty / `1` / `default` | Enabled in the platform default cache directory: `%LOCALAPPDATA%\llama.cpp\cl-cache` (Windows), `~/Library/Caches/llama.cpp/cl-cache` (macOS), `<temp dir>/llama.cpp/cl-cache` elsewhere. |
|
||||
| `0` / `off` / `none` / `disable(d)` | Disabled. |
|
||||
| any other value | Used verbatim as the cache directory path. |
|
||||
|
||||
If the chosen directory cannot be created or used, the cache disables itself for the process
|
||||
and kernels are compiled from source as usual. Set `GGML_OPENCL_KERNEL_CACHE_DEBUG=1` to
|
||||
print a HIT/MISS/SAVE trace to stderr.
|
||||
|
||||
## Android
|
||||
|
||||
Ubuntu 22.04 is used for targeting Android. Make sure the following tools are accessible from command line,
|
||||
|
||||
@@ -794,6 +794,8 @@ use 1 SYCL GPUs: [0] with Max compute units:512
|
||||
| GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
|
||||
| GGML_SYCL_USE_LEVEL_ZERO_API | 1 (default) or 0 | Use Level Zero API for device memory allocation instead of SYCL. Reduces system RAM usage on Intel dGPUs by avoiding DMA-buf/TTM host memory staging. Requires GGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON at build time. SYCL backend always runs on Level Zero running time even if it's set as OFF (The SYCL api will be usage for memory allocation).|
|
||||
| GGML_SYCL_ENABLE_DNN | 0 or 1 (default)| Enable running computations through oneDNN and always use oneMKL. |
|
||||
| GGML_SYCL_FA_ONEDNN | 1 (default) or 0 | Enable the oneDNN fused SDPA (flash-attention) path on supported GPUs. Set to 0 to always use the native SYCL flash-attention kernel. |
|
||||
| GGML_SYCL_FA_ONEDNN_MAX_KV | 0 (default, disabled) or positive integer | By default (0), all sequences are handled by the oneDNN fused SDPA path, regardless of KV length; a positive value caps that length, past which sequences fall back to the native kernel. If GPU driver watchdog resets (DEVICE_LOST) occur during long-context inference, set this near the context depth where they start, e.g. 24576. |
|
||||
| GGML_SYCL_ENABLE_VMM | 0 or 1 (default) | Enable the virtual-memory device pool. |
|
||||
| GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute (currently top-k MoE gating). |
|
||||
| ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.<br>Recommended to use when --split-mode = layer |
|
||||
|
||||
@@ -361,12 +361,6 @@ You can download it from your Linux distro's package manager or from here: [ROCm
|
||||
|
||||
Note: `GPU_TARGETS` is optional, omitting it will build the code for all GPUs in the current system.
|
||||
|
||||
To enhance flash attention performance on RDNA3+ or CDNA architectures, you can utilize the rocWMMA library by enabling the `-DGGML_HIP_ROCWMMA_FATTN=ON` option. This requires rocWMMA headers to be installed on the build system.
|
||||
|
||||
The rocWMMA library is included by default when installing the ROCm SDK using the `rocm` meta package provided by AMD. Alternatively, if you are not using the meta package, you can install the library using the `rocwmma-dev` or `rocwmma-devel` package, depending on your system's package manager.
|
||||
|
||||
As an alternative, you can manually install the library by cloning it from the official [GitHub repository](https://github.com/ROCm/rocWMMA), checkout the corresponding version tag (e.g. `rocm-6.2.4`) and set `-DCMAKE_CXX_FLAGS="-I<path/to/rocwmma>/library/include/"` in CMake. This also works under Windows despite not officially supported by AMD.
|
||||
|
||||
Note that if you get the following error:
|
||||
```
|
||||
clang: error: cannot find ROCm device library; provide its path via '--rocm-path' or '--rocm-device-lib-path', or pass '-nogpulib' to build without ROCm device library
|
||||
|
||||
@@ -45,6 +45,8 @@ class MyModel(MmprojModel):
|
||||
|
||||
Add an enum entry in `MODEL_ARCH`, the model human friendly name in `MODEL_ARCH_NAMES` and the GGUF tensor names in `MODEL_TENSORS`.
|
||||
|
||||
NOTE: Pick the GGUF arch string (and the matching `src/models/<name>.cpp` filename, see section 3) carefully up front, following existing naming conventions. Once GGUF files are published under a given arch string, renaming it later breaks the community's existing files, so this is not something to leave for cleanup in a follow-up PR.
|
||||
|
||||
Example for `falcon` model:
|
||||
```python
|
||||
MODEL_ARCH.FALCON: [
|
||||
@@ -101,6 +103,7 @@ The model params and tensors layout must be defined in `llama.cpp` source files:
|
||||
- You may also need to update `LLM_KV_NAMES`, `LLM_TENSOR_NAMES` and `LLM_TENSOR_INFOS`
|
||||
3. Add any non-standard metadata loading in the `llama_model_loader` constructor in `src/llama-model-loader.cpp`.
|
||||
4. If the model has a RoPE operation, add a case for the architecture in `llama_model_rope_type` function in `src/llama-model.cpp`.
|
||||
5. Check for other places that switch/iterate over every `llm_arch` value, e.g. `src/llama-model-saver.cpp` and any mandatory-hparam lists (such as which archs require MoE metadata). Grep for `LLM_ARCH_` usages to find them. Missing one of these is a common cause of CI test failures (e.g. `test-llama-archs`) after adding a new arch.
|
||||
|
||||
NOTE: The dimensions in `ggml` are typically in the reverse order of the `pytorch` dimensions.
|
||||
|
||||
@@ -133,6 +136,16 @@ Note:
|
||||
|
||||
## Tips and tricks
|
||||
|
||||
### Prefer conversion-time tensor modifications over graph-time ones
|
||||
|
||||
If the model contains constant modifications of tensors in the graph (for example, `norm(1 + weight)`) or performs tensor permutations/chunking, perform the modifications during conversion rather than in the graph code. This keeps the inference graph simpler and avoids extra runtime ops.
|
||||
|
||||
Examples:
|
||||
- Gemma 3 folds the `1 +` of its `norm(1 + weight)` normalization into the weights at conversion time, so the graph just does a plain RMS norm.
|
||||
- Qwen3-Next applies its tensor permutation during conversion (in `modify_tensors`), so the graph can consume the already-permuted weights directly.
|
||||
|
||||
Exception: a plain `weight * scale` with a constant scale is usually better left to inference time rather than folded into the weight at conversion. The scale conceptually applies to the activation, not the weight, so folding it into the weight can hurt numerical stability, and it shifts the weight's value range in a way that can make quantization worse. In this case, write the scale to GGUF as its own metadata key (e.g. `%s.attention.output_scale`, `%s.attention.value_scale`, `%s.embedding_scale`) and apply it in the graph, instead of pre-multiplying the weight tensor during conversion.
|
||||
|
||||
### Working with ggml_rope_ext
|
||||
|
||||
PyTorch implementations usually prefer explicitly calculating `freq_cis`/`sin`/`cos` components. However, in llama.cpp, most RoPE operations can be handled via `ggml_rope_ext`, which does not require a sin/cos matrix. This saves memory while allowing the GGML RoPE kernel to be fused with other ops.
|
||||
|
||||
+5
-5
@@ -16,22 +16,22 @@ conda-forge provides builds for:
|
||||
- Apple Metal (macOS)
|
||||
|
||||
```sh
|
||||
conda install -c conda-forge llama-cpp
|
||||
conda install -c conda-forge llama.cpp
|
||||
```
|
||||
|
||||
```sh
|
||||
mamba install -c conda-forge llama-cpp
|
||||
mamba install -c conda-forge llama.cpp
|
||||
```
|
||||
|
||||
```sh
|
||||
# Project-local installation
|
||||
pixi add llama-cpp
|
||||
pixi add llama.cpp
|
||||
|
||||
# Global installation
|
||||
pixi global install llama-cpp
|
||||
pixi global install llama.cpp
|
||||
```
|
||||
|
||||
This distribution is managed on [`conda-forge/llama-cpp-feedstock`](https://github.com/conda-forge/llama.cpp-feedstock/).
|
||||
This distribution is managed on [`conda-forge/llama.cpp-feedstock`](https://github.com/conda-forge/llama.cpp-feedstock/).
|
||||
|
||||
Shall you have any problems, please open an issue on [its issue tracker](https://github.com/conda-forge/llama.cpp-feedstock/issues).
|
||||
|
||||
|
||||
+34
-1
@@ -78,6 +78,38 @@ See:
|
||||
|
||||
- #22105
|
||||
|
||||
### DSpark (`draft-dspark`)
|
||||
|
||||
DSpark extends DFlash with a semi-autoregressive _Markov head_: the draft still emits a whole
|
||||
block per forward pass, but each block position's logits are biased by a low-rank term keyed on
|
||||
the previous token, chained in-graph across the block. This keeps drafting at one decode per
|
||||
block while recovering some of the left-to-right signal that pure block diffusion loses.
|
||||
|
||||
The draft is a small DeepSpec checkpoint trained for a specific target (for example
|
||||
[`deepseek-ai/dspark_qwen3_4b_block7`](https://huggingface.co/deepseek-ai/dspark_qwen3_4b_block7)
|
||||
for `Qwen/Qwen3-4B`). Convert it with `--target-model-dir` so it inherits the target's tokenizer
|
||||
and token embeddings:
|
||||
|
||||
```bash
|
||||
python convert_hf_to_gguf.py deepseek-ai/dspark_qwen3_4b_block7 \
|
||||
--target-model-dir Qwen/Qwen3-4B --outtype bf16 --outfile Qwen3-4B-DSpark.gguf
|
||||
|
||||
llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DSpark.gguf \
|
||||
--spec-type draft-dspark --spec-draft-n-max 7 -fa on --jinja
|
||||
```
|
||||
|
||||
`--spec-draft-n-max` is clamped to the draft model's trained block size.
|
||||
|
||||
`--spec-draft-conf-min P` truncates each drafted block at the first position whose predicted
|
||||
acceptance (from the draft's confidence head, if present) falls below `P` (default 0 = disabled).
|
||||
|
||||
Currently only drafts with a Qwen3 backbone are supported; support for other backbones
|
||||
(e.g. Gemma4) is planned.
|
||||
|
||||
See:
|
||||
|
||||
- #25173
|
||||
|
||||
### n-gram Cache (`ngram-cache`)
|
||||
|
||||
An n-gram is a sequence of n tokens. The n-gram cache implementation maintains statistics about short n-gram sequences.
|
||||
@@ -173,7 +205,7 @@ If a draft model is combined with a draftless decoding the draftless decoding ha
|
||||
### General Speculative Parameters
|
||||
|
||||
```
|
||||
--spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod]
|
||||
--spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-dspark|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod]
|
||||
comma-separated list of types of speculative decoding to use
|
||||
(default: none)
|
||||
(env: LLAMA_ARG_SPEC_TYPE)
|
||||
@@ -314,6 +346,7 @@ Specifies a comma-separated list of speculative decoding types to use.
|
||||
| `draft-simple` | Use a simple draft model for speculation |
|
||||
| `draft-eagle3` | Use an EAGLE-3 draft model that reads the target's hidden states |
|
||||
| `draft-dflash` | Use a DFlash block-diffusion draft model that emits a block per step |
|
||||
| `draft-dspark` | Use a DSpark draft model (DFlash backbone + semi-autoregressive Markov head) |
|
||||
| `draft-mtp` | Use Multi Token Prediction (MTP) heads from the main model |
|
||||
| `ngram-cache` | Use n-gram cache lookup |
|
||||
| `ngram-simple` | Use simple n-gram pattern matching |
|
||||
|
||||
@@ -216,7 +216,6 @@ option(GGML_HIP "ggml: use HIP"
|
||||
option(GGML_HIP_GRAPHS "ggml: use HIP graph" ON)
|
||||
option(GGML_HIP_RCCL "ggml: use ROCm Collective Comm. Library" OFF)
|
||||
option(GGML_HIP_NO_VMM "ggml: do not try to use HIP VMM" ON)
|
||||
option(GGML_HIP_ROCWMMA_FATTN "ggml: enable rocWMMA for FlashAttention" OFF)
|
||||
option(GGML_HIP_MMQ_MFMA "ggml: enable MFMA MMA for CDNA in MMQ" ON)
|
||||
option(GGML_HIP_EXPORT_METRICS "ggml: enable kernel perf metrics output" OFF)
|
||||
option(GGML_MUSA_GRAPHS "ggml: use MUSA graph, experimental, unstable" OFF)
|
||||
|
||||
+26
-17
@@ -906,26 +906,35 @@ static int ggml_backend_sched_backend_id_from_cur(ggml_backend_sched_t sched, st
|
||||
}
|
||||
|
||||
// operations with weights are preferably run on the same backend as the weights
|
||||
for (int i = 0; i < GGML_MAX_SRC; i++) {
|
||||
const struct ggml_tensor * src = tensor->src[i];
|
||||
if (src == NULL) {
|
||||
continue;
|
||||
}
|
||||
// skip ROPE since the rope freqs tensor is too small to choose a backend based on it
|
||||
// not an ideal solution
|
||||
if (tensor->op != GGML_OP_ROPE && src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
|
||||
int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor);
|
||||
// check if a backend with higher prio wants to offload the op
|
||||
if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) {
|
||||
for (int b = 0; b < src_backend_id; b++) {
|
||||
if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) {
|
||||
SET_CAUSE(tensor, "1.off");
|
||||
return b;
|
||||
// TODO: there are exceptions (see below) - not an ideal solution
|
||||
bool allow = true;
|
||||
|
||||
// skip ROPE since the rope freqs tensor is too small to choose a backend based on it
|
||||
allow = allow && tensor->op != GGML_OP_ROPE;
|
||||
|
||||
// skip FLASH_ATTN_EXT since the sinks tensor is too small to choose a based based on it
|
||||
allow = allow && tensor->op != GGML_OP_FLASH_ATTN_EXT;
|
||||
|
||||
if (allow) {
|
||||
for (int i = 0; i < GGML_MAX_SRC; i++) {
|
||||
const struct ggml_tensor * src = tensor->src[i];
|
||||
if (src == NULL) {
|
||||
continue;
|
||||
}
|
||||
if (src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
|
||||
int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor);
|
||||
// check if a backend with higher prio wants to offload the op
|
||||
if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) {
|
||||
for (int b = 0; b < src_backend_id; b++) {
|
||||
if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) {
|
||||
SET_CAUSE(tensor, "1.off");
|
||||
return b;
|
||||
}
|
||||
}
|
||||
}
|
||||
SET_CAUSE(tensor, "1.wgt%d", i);
|
||||
return src_backend_id;
|
||||
}
|
||||
SET_CAUSE(tensor, "1.wgt%d", i);
|
||||
return src_backend_id;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1797,14 +1797,6 @@ class tinyBLAS_Q0_AVX {
|
||||
//PPC Implementation
|
||||
#if defined(__MMA__)
|
||||
|
||||
#define SAVE_ACC(ACC, ii, jj) \
|
||||
__builtin_mma_disassemble_acc(vec_C, ACC); \
|
||||
for (int I = 0; I < 4; I++) { \
|
||||
for (int J = 0; J < 4; J++) { \
|
||||
*((float*)(C+ii+((jj+J)*ldc)+I)) = *((float*)&vec_C[I]+J); \
|
||||
} \
|
||||
} \
|
||||
|
||||
template<typename T>
|
||||
struct mma_instr;
|
||||
|
||||
@@ -1834,10 +1826,49 @@ class tinyBLAS_HP16_PPC {
|
||||
}
|
||||
|
||||
void matmul(int64_t m, int64_t n) {
|
||||
mnpack(0, m, 0, n);
|
||||
int64_t mc = 256;
|
||||
int64_t nc = 256;
|
||||
int64_t kc = 256;
|
||||
#if defined(_AIX) || defined(__BIG_ENDIAN__)
|
||||
mc = 128;
|
||||
nc = 128;
|
||||
kc = 128;
|
||||
#endif
|
||||
if (k < kc) {
|
||||
kc = k;
|
||||
}
|
||||
bool can_use_tiled = (m % mc == 0) && (n % nc == 0) && (k % kc == 0);
|
||||
if (can_use_tiled) {
|
||||
matmul_tiled(m, n, mc, nc, kc);
|
||||
} else {
|
||||
mnpack(0, m, 0, n);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
__attribute__((always_inline))
|
||||
inline void save_acc(acc_t * ACC, int64_t ii, int64_t jj) {
|
||||
vec_t vec_C[4];
|
||||
__builtin_mma_disassemble_acc(vec_C, ACC);
|
||||
for (int I = 0; I < 4; I++) {
|
||||
for (int J = 0; J < 4; J++) {
|
||||
*((float *)(C+ii+((jj+J)*ldc)+I)) = *((float *)&vec_C[I]+J);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__attribute__((always_inline))
|
||||
inline void add_save_acc(acc_t * ACC, int64_t ii, int64_t jj) {
|
||||
vec_t vec_C[4];
|
||||
__builtin_mma_disassemble_acc(vec_C, ACC);
|
||||
for (int I = 0; I < 4; I++) {
|
||||
for (int J = 0; J < 4; J++) {
|
||||
float * c_ptr = (float *)(C+ii+((jj+J)*ldc)+I);
|
||||
*c_ptr += *((float *)&vec_C[I]+J);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void vector_permute_store(vec_t *c, int numVec, unsigned char *vecOffset) {
|
||||
vec_t t[8], s[8];
|
||||
vec_t swiz1 = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23};
|
||||
@@ -1896,6 +1927,7 @@ class tinyBLAS_HP16_PPC {
|
||||
j = (rows >> 3);
|
||||
if (j > 0) {
|
||||
do {
|
||||
aoffsets[0] = aoffset;
|
||||
if (cols == 4) {
|
||||
aoffsets[0] = aoffset;
|
||||
for (int it = 1; it < 4; ++it)
|
||||
@@ -1910,17 +1942,17 @@ class tinyBLAS_HP16_PPC {
|
||||
}
|
||||
i = (cols >> 3);
|
||||
if (i > 0) {
|
||||
aoffsets[0] = aoffset;
|
||||
for (int it = 1; it < 8; ++it) {
|
||||
aoffsets[it] = aoffsets[it-1] + lda;
|
||||
}
|
||||
aoffset += 8 * lda;
|
||||
|
||||
do {
|
||||
for (int it = 0; it < 8; ++it)
|
||||
c_arr[it] = vec_xl(0, (vector unsigned char*)aoffsets[it]);
|
||||
vector_permute_store(c_arr, 8, vecOffset);
|
||||
for (int it = 0; it < 8; ++it)
|
||||
aoffsets[it] = aoffsets[it] + 8*lda;
|
||||
aoffsets[it] = aoffsets[it] + 8;
|
||||
vecOffset += 128;
|
||||
i--;
|
||||
} while(i > 0);
|
||||
@@ -2147,8 +2179,8 @@ class tinyBLAS_HP16_PPC {
|
||||
mma_instr<TA>::outer_product(&acc_1, vec_A[x], vec_B[x+4]);
|
||||
}
|
||||
}
|
||||
SAVE_ACC(&acc_0, ii, jj);
|
||||
SAVE_ACC(&acc_1, ii, jj+4);
|
||||
save_acc(&acc_0, ii, jj);
|
||||
save_acc(&acc_1, ii, jj+4);
|
||||
}
|
||||
|
||||
void KERNEL_8x4(int64_t ii, int64_t jj) {
|
||||
@@ -2164,8 +2196,8 @@ class tinyBLAS_HP16_PPC {
|
||||
mma_instr<TA>::outer_product(&acc_1, vec_A[x+4], vec_B[x]);
|
||||
}
|
||||
}
|
||||
SAVE_ACC(&acc_0, ii, jj);
|
||||
SAVE_ACC(&acc_1, ii+4, jj);
|
||||
save_acc(&acc_0, ii, jj);
|
||||
save_acc(&acc_1, ii+4, jj);
|
||||
}
|
||||
|
||||
|
||||
@@ -2186,13 +2218,64 @@ class tinyBLAS_HP16_PPC {
|
||||
mma_instr<TA>::outer_product(&acc_3, vec_A[x+4], vec_B[x+4]);
|
||||
}
|
||||
}
|
||||
|
||||
SAVE_ACC(&acc_0, ii, jj);
|
||||
SAVE_ACC(&acc_1, ii, jj+4);
|
||||
SAVE_ACC(&acc_2, ii+4, jj);
|
||||
SAVE_ACC(&acc_3, ii+4, jj+4);
|
||||
save_acc(&acc_0, ii, jj);
|
||||
save_acc(&acc_1, ii, jj+4);
|
||||
save_acc(&acc_2, ii+4, jj);
|
||||
save_acc(&acc_3, ii+4, jj+4);
|
||||
}
|
||||
|
||||
inline void MMA_16x8(vec_t * vec_A0, vec_t * vec_A1, vec_t * vec_B, acc_t * acc) {
|
||||
for (int x = 0; x < 4; x ++) {
|
||||
mma_instr<TA>::outer_product(&acc[0], vec_A0[x], vec_B[x]);
|
||||
mma_instr<TA>::outer_product(&acc[1], vec_A0[x], vec_B[x+4]);
|
||||
mma_instr<TA>::outer_product(&acc[2], vec_A0[x+4], vec_B[x]);
|
||||
mma_instr<TA>::outer_product(&acc[3], vec_A0[x+4], vec_B[x+4]);
|
||||
mma_instr<TA>::outer_product(&acc[4], vec_A1[x], vec_B[x]);
|
||||
mma_instr<TA>::outer_product(&acc[5], vec_A1[x], vec_B[x+4]);
|
||||
mma_instr<TA>::outer_product(&acc[6], vec_A1[x+4], vec_B[x]);
|
||||
mma_instr<TA>::outer_product(&acc[7], vec_A1[x+4], vec_B[x+4]);
|
||||
}
|
||||
}
|
||||
void KERNEL(int64_t ii, int64_t jj, int64_t mc, int64_t nc, int64_t kc, vec_t * vec_A, vec_t * vec_B, int64_t kk) {
|
||||
for (int64_t i = 0; i < mc; i += 16) {
|
||||
int A_base_addr = (mc / 8) * (i / 8) * 8;
|
||||
for (int64_t j = 0; j < nc; j += 8) {
|
||||
int B_base_addr = (nc / 8) * (j / 8) * 8;
|
||||
acc_t acc[8];
|
||||
vec_t A0_block[8]; vec_t A1_block[8];
|
||||
for (int x = 0; x < 8; x++)
|
||||
__builtin_mma_xxsetaccz(&acc[x]);
|
||||
for (int64_t l = 0; l < kc; l += 8) {
|
||||
int A0_block_idx = A_base_addr + (l / 8) * 8;
|
||||
int A1_block_idx = A0_block_idx + (mc / 8) * 8;
|
||||
int B_block_idx = B_base_addr + (l / 8) * 8;
|
||||
vec_t* A0_block = &vec_A[A0_block_idx];
|
||||
vec_t* A1_block = &vec_A[A1_block_idx];
|
||||
vec_t* B_block = &vec_B[B_block_idx];
|
||||
MMA_16x8(A0_block, A1_block, B_block, acc);
|
||||
}
|
||||
if (kk == 0) {
|
||||
save_acc(&acc[0], ii + i, jj + j);
|
||||
save_acc(&acc[1], ii + i, jj + j + 4);
|
||||
save_acc(&acc[2], ii + i + 4, jj + j);
|
||||
save_acc(&acc[3], ii + i + 4, jj + j + 4);
|
||||
save_acc(&acc[4], ii + i + 8, jj + j);
|
||||
save_acc(&acc[5], ii + i + 8, jj + j + 4);
|
||||
save_acc(&acc[6], ii + i + 12, jj + j);
|
||||
save_acc(&acc[7], ii + i + 12, jj + j + 4);
|
||||
} else {
|
||||
add_save_acc(&acc[0], ii + i, jj + j);
|
||||
add_save_acc(&acc[1], ii + i, jj + j + 4);
|
||||
add_save_acc(&acc[2], ii + i + 4, jj + j);
|
||||
add_save_acc(&acc[3], ii + i + 4, jj + j + 4);
|
||||
add_save_acc(&acc[4], ii + i + 8, jj + j);
|
||||
add_save_acc(&acc[5], ii + i + 8, jj + j + 4);
|
||||
add_save_acc(&acc[6], ii + i + 12, jj + j);
|
||||
add_save_acc(&acc[7], ii + i + 12, jj + j + 4);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
template<int RM, int RN>
|
||||
void gemm_small(int64_t m0, int64_t m, int64_t n0, int64_t n) {
|
||||
int64_t ytiles = (m - m0) / RM;
|
||||
@@ -2281,6 +2364,29 @@ class tinyBLAS_HP16_PPC {
|
||||
}
|
||||
}
|
||||
|
||||
void matmul_tiled(int64_t m, int64_t n, int64_t mc, int64_t nc, int64_t kc) {
|
||||
int64_t ytiles = m / mc;
|
||||
int64_t xtiles = n / nc;
|
||||
int64_t tiles = xtiles * ytiles;
|
||||
int64_t duty = (tiles + nth - 1) / nth;
|
||||
int64_t start = duty * ith;
|
||||
int64_t end = start + duty;
|
||||
if (end > tiles) {
|
||||
end = tiles;
|
||||
}
|
||||
for (int64_t job = start; job < end; ++job) {
|
||||
int64_t ii = (job / xtiles) * mc;
|
||||
int64_t jj = (job % xtiles) * nc;
|
||||
for (int64_t kk = 0; kk < k; kk += kc) {
|
||||
vec_t A_pack[kc * mc / 8];
|
||||
vec_t B_pack[kc * nc / 8];
|
||||
packNormal(A + (ii * lda) + kk, lda, kc, mc, (uint8_t *)A_pack);
|
||||
packNormal(B + (jj * ldb) + kk, ldb, kc, nc, (uint8_t *)B_pack);
|
||||
KERNEL(ii, jj, mc, nc, kc, A_pack, B_pack, kk);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int RM, int RN>
|
||||
NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) {
|
||||
int64_t ytiles = (m - m0) / RM;
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
#include "common.cuh"
|
||||
#include "fattn-tile.cuh"
|
||||
#include "fattn-wmma-f16.cuh"
|
||||
|
||||
void ggml_cuda_flash_attn_ext_tile(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
||||
const ggml_tensor * K = dst->src[1];
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
#include "common.cuh"
|
||||
#include "fattn-common.cuh"
|
||||
#include "fattn-wmma-f16.cuh"
|
||||
|
||||
// nbatch_fa == number of KQ rows to process per iteration
|
||||
// nbatch_K == number of K columns to load in parallel for KQ calculation
|
||||
@@ -825,12 +824,7 @@ static __global__ void flash_attn_tile(
|
||||
|
||||
// Skip unused kernel variants for faster compilation:
|
||||
|
||||
if (
|
||||
#ifdef GGML_USE_WMMA_FATTN
|
||||
(ncols2 != 1 && DV != 40 && DV != 72 && DV != 512) ||
|
||||
#endif // GGML_USE_WMMA_FATTN
|
||||
(use_logit_softcap && !(DV == 128 || DV == 256 || DV == 512))
|
||||
) {
|
||||
if ((use_logit_softcap && !(DV == 128 || DV == 256 || DV == 512))) {
|
||||
GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale,
|
||||
max_bias, m0, m1, n_head_log2, logit_softcap,
|
||||
ne00, ne01, ne02, ne03,
|
||||
|
||||
@@ -1,705 +0,0 @@
|
||||
// Old and deprecated WMMA FlashAttention implementation.
|
||||
// It is still needed for Volta since the memory layout of NVIDIA tensor cores changed with Turing.
|
||||
// Long-term the WMMA code should be replaced with a dedicated Volta implementation.
|
||||
|
||||
#include "common.cuh"
|
||||
#include "fattn-common.cuh"
|
||||
#include "fattn-wmma-f16.cuh"
|
||||
|
||||
#ifdef GGML_USE_WMMA_FATTN
|
||||
#if !defined(GGML_USE_HIP)
|
||||
#include <mma.h>
|
||||
#if defined(GGML_USE_MUSA)
|
||||
namespace wmma = mtmusa::wmma;
|
||||
#else // GGML_USE_MUSA
|
||||
namespace wmma = nvcuda::wmma;
|
||||
#endif // GGML_USE_MUSA
|
||||
#elif defined(GGML_USE_HIP)
|
||||
#include <rocwmma/rocwmma.hpp>
|
||||
namespace wmma = rocwmma;
|
||||
#endif // !defined(GGML_USE_HIP)
|
||||
#endif // GGML_USE_WMMA_FATTN
|
||||
|
||||
// D == head size, VKQ_stride == num VKQ rows calculated in parallel:
|
||||
template<int D, int ncols, int nwarps, int VKQ_stride, typename KQ_acc_t, bool use_logit_softcap>
|
||||
__launch_bounds__(nwarps*ggml_cuda_get_physical_warp_size(), 1)
|
||||
static __global__ void flash_attn_ext_f16(
|
||||
const char * Q_ptr,
|
||||
const char * K_ptr,
|
||||
const char * V_ptr,
|
||||
const char * mask_ptr,
|
||||
const char * sinks_ptr,
|
||||
const int * KV_max_ptr,
|
||||
float * dst_ptr,
|
||||
float2 * dst_meta_ptr,
|
||||
const float scale,
|
||||
const float max_bias,
|
||||
const float m0,
|
||||
const float m1,
|
||||
const uint32_t n_head_log2,
|
||||
const float logit_softcap,
|
||||
const int32_t ne00, const uint3 ne01, const int32_t ne02, const int32_t ne03,
|
||||
const int32_t nb01, const int32_t nb02, const int32_t nb03,
|
||||
const int32_t ne10, const int32_t ne11, const int32_t ne12, const int32_t ne13,
|
||||
const int32_t nb11, const int32_t nb12, const int64_t nb13,
|
||||
const int32_t nb21, const int32_t nb22, const int64_t nb23,
|
||||
const int32_t ne31, const int32_t ne32, const int32_t ne33,
|
||||
const int32_t nb31, const int32_t nb32, const int64_t nb33) {
|
||||
#if defined(FLASH_ATTN_AVAILABLE) && (defined(GGML_HIP_ROCWMMA_FATTN) && defined(GGML_USE_WMMA_FATTN))
|
||||
const char * GGML_CUDA_RESTRICT Q = Q_ptr;
|
||||
const char * GGML_CUDA_RESTRICT K = K_ptr;
|
||||
const char * GGML_CUDA_RESTRICT V = V_ptr;
|
||||
const char * GGML_CUDA_RESTRICT mask = mask_ptr;
|
||||
const char * GGML_CUDA_RESTRICT sinks = sinks_ptr;
|
||||
const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr;
|
||||
float * GGML_CUDA_RESTRICT dst = dst_ptr;
|
||||
float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
|
||||
// Skip unused kernel variants for faster compilation:
|
||||
if (use_logit_softcap && !(D == 128 || D == 256)) {
|
||||
NO_DEVICE_CODE;
|
||||
return;
|
||||
}
|
||||
|
||||
//In this kernel Q, K, V are matrices while i, j, k are matrix indices.
|
||||
|
||||
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
|
||||
|
||||
const int ic0 = ncols*blockIdx.x; // Index of the first Q/QKV column to work on.
|
||||
|
||||
static_assert(D <= FATTN_KQ_STRIDE, "D must be <= FATTN_KQ_STRIDE.");
|
||||
static_assert(ncols == 8 || ncols % 16 == 0, "ncols must be 8 or a multiple of 16.");
|
||||
constexpr int frag_m = ncols == 8 ? 32 : 16;
|
||||
constexpr int frag_n = ncols == 8 ? 8 : 16;
|
||||
static_assert(D % frag_m == 0, "If ncols == 8 then D % frag_m must be 0.");
|
||||
#if defined(GGML_USE_HIP) && HIP_VERSION >= 60500000
|
||||
typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, _Float16, wmma::row_major> frag_a_K;
|
||||
typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, _Float16, wmma::col_major> frag_a_V;
|
||||
typedef wmma::fragment<wmma::matrix_b, frag_m, frag_n, 16, _Float16, wmma::col_major> frag_b;
|
||||
typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, KQ_acc_t> frag_c_KQ;
|
||||
typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, _Float16> frag_c_VKQ;
|
||||
#else
|
||||
typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, half, wmma::row_major> frag_a_K;
|
||||
typedef wmma::fragment<wmma::matrix_a, frag_m, frag_n, 16, half, wmma::col_major> frag_a_V;
|
||||
typedef wmma::fragment<wmma::matrix_b, frag_m, frag_n, 16, half, wmma::col_major> frag_b;
|
||||
typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, KQ_acc_t> frag_c_KQ;
|
||||
typedef wmma::fragment<wmma::accumulator, frag_m, frag_n, 16, half> frag_c_VKQ;
|
||||
#endif
|
||||
|
||||
constexpr int KQ_stride_tc = nwarps*frag_m; // Number of KQ rows calculated in parallel.
|
||||
constexpr int VKQ_ratio = KQ_stride_tc/VKQ_stride; // Number of parallel VKQ accumulators needed to keep all warps busy.
|
||||
static_assert(VKQ_ratio <= nwarps, "VKQ_ratio must be <= nwarps.");
|
||||
|
||||
// Pad internal representation of KQ, KQV to reduce shared memory bank conflicts:
|
||||
constexpr int D_padded = D + 8;
|
||||
constexpr int kqs_padded = FATTN_KQ_STRIDE + 8;
|
||||
constexpr int kqar = sizeof(KQ_acc_t)/sizeof(half);
|
||||
|
||||
ggml_cuda_pdl_sync();
|
||||
const int sequence = blockIdx.z / ne02;
|
||||
const int head = blockIdx.z - sequence*ne02;
|
||||
const int gqa_ratio = ne02 / ne12; // With grouped query attention there are > 1 Q matrices per K, V matrix.
|
||||
const float * Q_f = (const float *) (Q + nb03* sequence + nb02* head + nb01*ic0);
|
||||
const half * K_h = (const half *) (K + nb13* sequence + nb12*(head / gqa_ratio));
|
||||
const half * V_h = (const half *) (V + nb13* sequence + nb12*(head / gqa_ratio)); // K and V have same shape
|
||||
const half * maskh = (const half *) (mask + nb33*(sequence % ne33) + nb31*ic0);
|
||||
const half2 * mask2 = (const half2 *) maskh;
|
||||
const float * sinksf = (const float *) sinks;
|
||||
|
||||
const int stride_Q = nb01 / sizeof(float);
|
||||
const int stride_KV = nb11 / sizeof(half);
|
||||
|
||||
const float slopef = get_alibi_slope(max_bias, head, n_head_log2, m0, m1);
|
||||
const half slopeh = __float2half(slopef);
|
||||
const half2 slope2 = make_half2(slopef, slopef);
|
||||
|
||||
const half2 logit_softcap_2 = make_half2(logit_softcap, logit_softcap);
|
||||
|
||||
frag_b Q_b[D/16][ncols/frag_n];
|
||||
|
||||
// A single buffer for temporarily holding tiles of KQ and VKQ parts:
|
||||
constexpr int mem_KQ = ncols*kqs_padded*kqar;
|
||||
constexpr int mem_VKQ_parts = VKQ_ratio*ncols*D_padded;
|
||||
__shared__ half KQ[mem_KQ >= mem_VKQ_parts ? mem_KQ : mem_VKQ_parts];
|
||||
float * KQ_f = (float *) KQ;
|
||||
half2 * KQ2 = (half2 *) KQ;
|
||||
|
||||
float KQ_rowsum_f[ncols/nwarps] = {0.0f};
|
||||
float KQ_max_f[ncols/nwarps];
|
||||
float KQ_max_scale_f[ncols/nwarps] = {0.0f};
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < ncols/nwarps; ++j) {
|
||||
KQ_max_f[j] = -FLT_MAX/2.0f;
|
||||
}
|
||||
|
||||
half2 KQ_rowsum_h2[ncols/nwarps] = {{0.0f, 0.0f}};
|
||||
half2 KQ_max_h2[ncols/nwarps];
|
||||
half2 KQ_max_scale_h2[ncols/nwarps] = {{0.0f, 0.0f}};
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < ncols/nwarps; ++j) {
|
||||
KQ_max_h2[j] = make_half2(-HALF_MAX_HALF, -HALF_MAX_HALF);
|
||||
}
|
||||
|
||||
__shared__ half VKQ[ncols*D_padded]; // Accumulator for final VKQ slice.
|
||||
half2 * VKQ2 = (half2 *) VKQ;
|
||||
|
||||
#if defined(GGML_USE_HIP) && HIP_VERSION >= 60500000
|
||||
const _Float16 * K_h_f16 = reinterpret_cast<const _Float16 *>(K_h);
|
||||
const _Float16 * V_h_f16 = reinterpret_cast<const _Float16 *>(V_h);
|
||||
_Float16 * KQ_f16 = reinterpret_cast<_Float16 *>(KQ);
|
||||
_Float16 * VKQ_f16 = reinterpret_cast<_Float16 *>(VKQ);
|
||||
#else
|
||||
const half * K_h_f16 = K_h;
|
||||
const half * V_h_f16 = V_h;
|
||||
half * KQ_f16 = KQ;
|
||||
half * VKQ_f16 = VKQ;
|
||||
#endif
|
||||
|
||||
#pragma unroll
|
||||
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
|
||||
const int j = j0 + threadIdx.y;
|
||||
#pragma unroll
|
||||
for (int i0 = 0; i0 < D/2; i0 += warp_size) {
|
||||
const int i = i0 + threadIdx.x;
|
||||
if (i0 + warp_size > D/2 && i >= D/2) {
|
||||
break;
|
||||
}
|
||||
VKQ2[j*(D_padded/2) + i] = make_half2(0.0f, 0.0f);
|
||||
}
|
||||
}
|
||||
|
||||
// Convert Q to half and apply scale, temporarily store in KQ:
|
||||
#pragma unroll
|
||||
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
|
||||
const int j = j0 + threadIdx.y;
|
||||
#pragma unroll
|
||||
for (int i0 = 0; i0 < D; i0 += warp_size) {
|
||||
const int i = i0 + threadIdx.x;
|
||||
if (i0 + warp_size > D && i >= D) {
|
||||
break;
|
||||
}
|
||||
KQ[j*D_padded + i] = ic0 + j < int(ne01.z) ? Q_f[j*stride_Q + i] * scale : 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Load Q into tensor core fragments/registers since it will be used frequently:
|
||||
#pragma unroll
|
||||
for (int i0 = 0; i0 < D; i0 += 16) {
|
||||
#pragma unroll
|
||||
for (int j0 = 0; j0 < ncols; j0 += frag_n) {
|
||||
wmma::load_matrix_sync(Q_b[i0/16][j0/frag_n], KQ_f16 + j0*D_padded + i0, D_padded);
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Iterate over ne11 == previous tokens:
|
||||
const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11;
|
||||
for (int k_VKQ_0 = blockIdx.y*FATTN_KQ_STRIDE; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*FATTN_KQ_STRIDE) {
|
||||
// Calculate tile of KQ:
|
||||
#pragma unroll
|
||||
for (int i_KQ_0 = 0; i_KQ_0 < FATTN_KQ_STRIDE; i_KQ_0 += KQ_stride_tc) {
|
||||
frag_c_KQ KQ_c[ncols/frag_n];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < ncols/frag_n; ++j) {
|
||||
wmma::fill_fragment(KQ_c[j], static_cast<KQ_acc_t>(0.0f));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int k_KQ_0 = 0; k_KQ_0 < D; k_KQ_0 += 16) {
|
||||
frag_a_K K_a;
|
||||
wmma::load_matrix_sync(K_a, K_h_f16 + int64_t(k_VKQ_0 + i_KQ_0 + frag_m*threadIdx.y)*stride_KV + k_KQ_0, stride_KV);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < ncols/frag_n; ++j) {
|
||||
wmma::mma_sync(KQ_c[j], K_a, Q_b[k_KQ_0/16][j], KQ_c[j]);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j0 = 0; j0 < ncols; j0 += frag_n) {
|
||||
wmma::store_matrix_sync((KQ_acc_t *) KQ + j0*kqs_padded + i_KQ_0 + frag_m*threadIdx.y, KQ_c[j0/frag_n], kqs_padded, wmma::mem_col_major);
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Calculate softmax for each KQ column using the current max. value.
|
||||
// The divisor is stored in KQ_rowsum and will be applied at the end.
|
||||
#pragma unroll
|
||||
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
|
||||
const int j = j0 + threadIdx.y;
|
||||
|
||||
if (std::is_same<KQ_acc_t, float>::value) {
|
||||
float KQ_f_tmp[FATTN_KQ_STRIDE / warp_size];
|
||||
#pragma unroll
|
||||
for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += warp_size) {
|
||||
const int k = k0 + threadIdx.x;
|
||||
|
||||
KQ_f_tmp[k0/warp_size] = KQ_f[j*kqs_padded + k];
|
||||
|
||||
if (use_logit_softcap) {
|
||||
KQ_f_tmp[k0/warp_size] = logit_softcap*tanhf(KQ_f_tmp[k0/warp_size]);
|
||||
}
|
||||
}
|
||||
|
||||
float KQ_max_new = KQ_max_f[j0/nwarps];
|
||||
#pragma unroll
|
||||
for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += warp_size) {
|
||||
const int k = k0 + threadIdx.x;
|
||||
|
||||
KQ_f_tmp[k0/warp_size] += mask && ic0 + j < int(ne01.z) ?
|
||||
__half2float(slopeh*maskh[j*(nb31/sizeof(half)) + k_VKQ_0 + k]) : 0.0f;
|
||||
KQ_max_new = max(KQ_max_new, KQ_f_tmp[k0/warp_size] + FATTN_KQ_MAX_OFFSET);
|
||||
}
|
||||
KQ_max_new = warp_reduce_max<warp_size>(KQ_max_new);
|
||||
|
||||
const float diff = KQ_max_f[j0/nwarps] - KQ_max_new;
|
||||
KQ_max_scale_f[j0/nwarps] = expf(diff);
|
||||
if (diff <= SOFTMAX_FTZ_THRESHOLD) {
|
||||
KQ_max_scale_f[j0/nwarps] = 0.0f;
|
||||
}
|
||||
KQ_max_f[j0/nwarps] = KQ_max_new;
|
||||
|
||||
float KQ_rowsum_add = 0.0f;
|
||||
#pragma unroll
|
||||
for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += warp_size) {
|
||||
const int k = k0 + threadIdx.x;
|
||||
|
||||
const float diff = KQ_f_tmp[k0/warp_size] - KQ_max_f[j0/nwarps];
|
||||
KQ_f_tmp[k0/warp_size] = expf(diff);
|
||||
if (diff <= SOFTMAX_FTZ_THRESHOLD) {
|
||||
KQ_f_tmp[k0/warp_size] = 0.0f;
|
||||
}
|
||||
KQ_rowsum_add += KQ_f_tmp[k0/warp_size];
|
||||
KQ[j*(kqar*kqs_padded) + k] = KQ_f_tmp[k0/warp_size];
|
||||
}
|
||||
KQ_rowsum_add = warp_reduce_sum<warp_size>(KQ_rowsum_add);
|
||||
|
||||
// Scale previous KQ_rowsum to account for a potential increase in KQ_max:
|
||||
KQ_rowsum_f[j0/nwarps] = KQ_max_scale_f[j0/nwarps]*KQ_rowsum_f[j0/nwarps] + KQ_rowsum_add;
|
||||
} else {
|
||||
half2 KQ2_tmp[FATTN_KQ_STRIDE/(2*warp_size)];
|
||||
#pragma unroll
|
||||
for (int k0 = 0; k0 < FATTN_KQ_STRIDE/2; k0 += warp_size) {
|
||||
const int k = k0 + threadIdx.x;
|
||||
|
||||
KQ2_tmp[k0/warp_size] = KQ2[j*(kqs_padded/2) + k];
|
||||
|
||||
if (use_logit_softcap) {
|
||||
// There is no dedicated tangens hyperbolicus function for half2.
|
||||
KQ2_tmp[k0/warp_size] = h2exp(KQ2_tmp[k0/warp_size]*make_half2(2.0f, 2.0f));
|
||||
KQ2_tmp[k0/warp_size] = (KQ2_tmp[k0/warp_size] - make_half2(1.0f, 1.0f))
|
||||
/(KQ2_tmp[k0/warp_size] + make_half2(1.0f, 1.0f));
|
||||
|
||||
KQ2_tmp[k0/warp_size] *= logit_softcap_2;
|
||||
}
|
||||
}
|
||||
|
||||
half2 KQ_max_new = KQ_max_h2[j0/nwarps];
|
||||
#pragma unroll
|
||||
for (int k0 = 0; k0 < FATTN_KQ_STRIDE/2; k0 += warp_size) {
|
||||
const int k = k0 + threadIdx.x;
|
||||
|
||||
KQ2_tmp[k0/warp_size] += mask && ic0 + j < int(ne01.z) ? slope2*mask2[(j*ne11 + k_VKQ_0)/2 + k] : make_half2(0.0f, 0.0f);
|
||||
KQ_max_new = ggml_cuda_hmax2(KQ_max_new, KQ2_tmp[k0/warp_size]);
|
||||
}
|
||||
KQ_max_new = __half2half2(warp_reduce_max<warp_size>(ggml_cuda_hmax(__low2half(KQ_max_new), __high2half(KQ_max_new))));
|
||||
const half2 diff = KQ_max_h2[j0/nwarps] - KQ_max_new;
|
||||
KQ_max_scale_h2[j0/nwarps] = h2exp(diff);
|
||||
const uint32_t ftz_mask = __hgt2_mask(diff, make_half2(SOFTMAX_FTZ_THRESHOLD, SOFTMAX_FTZ_THRESHOLD));
|
||||
*((uint32_t *) &KQ_max_scale_h2[j0/nwarps]) &= ftz_mask;
|
||||
KQ_max_h2[j0/nwarps] = KQ_max_new;
|
||||
|
||||
half2 KQ_rowsum_add = make_half2(0.0f, 0.0f);
|
||||
#pragma unroll
|
||||
for (int k0 = 0; k0 < FATTN_KQ_STRIDE/2; k0 += warp_size) {
|
||||
const int k = k0 + threadIdx.x;
|
||||
|
||||
const half2 diff = KQ2_tmp[k0/warp_size] - KQ_max_h2[j0/nwarps];
|
||||
KQ2_tmp[k0/warp_size] = h2exp(diff);
|
||||
const uint32_t ftz_mask = __hgt2_mask(diff, make_half2(SOFTMAX_FTZ_THRESHOLD, SOFTMAX_FTZ_THRESHOLD));
|
||||
*((uint32_t *) &KQ2_tmp[k0/warp_size]) &= ftz_mask;
|
||||
KQ_rowsum_add += KQ2_tmp[k0/warp_size];
|
||||
KQ2[j*(kqs_padded/2) + k] = KQ2_tmp[k0/warp_size];
|
||||
}
|
||||
KQ_rowsum_add = warp_reduce_sum<warp_size>(KQ_rowsum_add);
|
||||
|
||||
// Scale previous KQ_rowsum to account for a potential increase in KQ_max:
|
||||
KQ_rowsum_h2[j0/nwarps] = KQ_max_scale_h2[j0/nwarps]*KQ_rowsum_h2[j0/nwarps] + KQ_rowsum_add;
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
frag_b KQ_b[FATTN_KQ_STRIDE/(VKQ_ratio*16)][ncols/frag_n];
|
||||
#pragma unroll
|
||||
for (int j0 = 0; j0 < ncols; j0 += frag_n) {
|
||||
#pragma unroll
|
||||
for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += VKQ_ratio*16) {
|
||||
const int k = k0 + (threadIdx.y % VKQ_ratio)*16;
|
||||
wmma::load_matrix_sync(
|
||||
KQ_b[k0/(VKQ_ratio*16)][j0/frag_n],
|
||||
KQ_f16 + j0*(kqar*kqs_padded) + k,
|
||||
kqar*kqs_padded);
|
||||
}
|
||||
}
|
||||
|
||||
frag_c_VKQ VKQ_c[D/VKQ_stride][ncols/frag_n];
|
||||
#pragma unroll
|
||||
for (int i_VKQ_0 = 0; i_VKQ_0 < D; i_VKQ_0 += VKQ_stride) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < ncols/frag_n; ++j) {
|
||||
wmma::fill_fragment(VKQ_c[i_VKQ_0/VKQ_stride][j], static_cast<half>(0.0f));
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k0 = 0; k0 < FATTN_KQ_STRIDE; k0 += VKQ_ratio*16) {
|
||||
const int k = k0 + (threadIdx.y % VKQ_ratio)*16;
|
||||
|
||||
frag_a_V v_a;
|
||||
wmma::load_matrix_sync(v_a, V_h_f16 + int64_t(k_VKQ_0 + k)*stride_KV + i_VKQ_0 + frag_m*(threadIdx.y/VKQ_ratio), stride_KV);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < ncols/frag_n; ++j) {
|
||||
wmma::mma_sync(VKQ_c[i_VKQ_0/VKQ_stride][j], v_a, KQ_b[k0/(VKQ_ratio*16)][j], VKQ_c[i_VKQ_0/VKQ_stride][j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
const int offset_k = (threadIdx.y % VKQ_ratio) * (ncols*D_padded);
|
||||
#pragma unroll
|
||||
for (int i_KQ_0 = 0; i_KQ_0 < D; i_KQ_0 += VKQ_stride) {
|
||||
#pragma unroll
|
||||
for (int j0 = 0; j0 < ncols; j0 += frag_n) {
|
||||
wmma::store_matrix_sync(
|
||||
KQ_f16 + offset_k + j0*D_padded + i_KQ_0 + frag_m*(threadIdx.y/VKQ_ratio),
|
||||
VKQ_c[i_KQ_0/VKQ_stride][j0/frag_n],
|
||||
D_padded, wmma::mem_col_major);
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
#pragma unroll
|
||||
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
|
||||
const int j = j0 + threadIdx.y;
|
||||
|
||||
half2 VKQ_scale;
|
||||
if (std::is_same<KQ_acc_t, float>::value) {
|
||||
VKQ_scale = make_half2(KQ_max_scale_f[j0/nwarps], KQ_max_scale_f[j0/nwarps]);
|
||||
} else {
|
||||
VKQ_scale = KQ_max_scale_h2[j0/nwarps];
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i0 = 0; i0 < D/2; i0 += warp_size) {
|
||||
const int i = i0 + threadIdx.x;
|
||||
if (i0 + warp_size > D/2 && i >= D/2) {
|
||||
break;
|
||||
}
|
||||
|
||||
half2 VKQ_add = make_half2(0.0f, 0.0f);
|
||||
#pragma unroll
|
||||
for (int l = 0; l < VKQ_ratio; ++l) {
|
||||
VKQ_add += KQ2[l*(ncols*D_padded/2) + j*(D_padded/2) + i];
|
||||
}
|
||||
VKQ2[j*(D_padded/2) + i] = VKQ_scale*VKQ2[j*(D_padded/2) + i] + VKQ_add;
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Apply attention sinks
|
||||
if (sinksf && blockIdx.y == 0) {
|
||||
const float sinkf = sinksf[head];
|
||||
const half sinkh = __float2half(sinkf);
|
||||
|
||||
#pragma unroll
|
||||
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
|
||||
const int j = j0 + threadIdx.y;
|
||||
|
||||
if (std::is_same<KQ_acc_t, float>::value) {
|
||||
float kqmax_new = fmaxf(KQ_max_f[j0/nwarps], sinkf);
|
||||
|
||||
const float KQ_max_scale = expf(KQ_max_f[j0/nwarps] - kqmax_new);
|
||||
KQ_max_f[j0/nwarps] = kqmax_new;
|
||||
|
||||
KQ_rowsum_f[j0/nwarps] = KQ_rowsum_f[j0/nwarps] * KQ_max_scale + expf(sinkf - KQ_max_f[j0/nwarps]);
|
||||
|
||||
const half2 scale_h2 = make_half2(KQ_max_scale, KQ_max_scale);
|
||||
#pragma unroll
|
||||
for (int i0 = 0; i0 < D/2; i0 += warp_size) {
|
||||
const int i = i0 + threadIdx.x;
|
||||
if (i0 + warp_size > D/2 && i >= D/2) break;
|
||||
VKQ2[j*(D_padded/2) + i] *= scale_h2;
|
||||
}
|
||||
} else {
|
||||
half kqmax_old = __low2half(KQ_max_h2[j0/nwarps]);
|
||||
half kqmax_new = fmaxf(kqmax_old, sinkh);
|
||||
KQ_max_h2[j0/nwarps] = __half2half2(kqmax_new);
|
||||
|
||||
const half KQ_max_scale_h = hexp(kqmax_old - kqmax_new);
|
||||
const half2 KQ_max_scale = __half2half2(KQ_max_scale_h);
|
||||
|
||||
KQ_rowsum_h2[j0/nwarps] = KQ_rowsum_h2[j0/nwarps] * KQ_max_scale;
|
||||
const half val = hexp(sinkh - kqmax_new);
|
||||
KQ_rowsum_h2[j0/nwarps].x = __hadd(KQ_rowsum_h2[j0/nwarps].x, val);
|
||||
|
||||
#pragma unroll
|
||||
for (int i0 = 0; i0 < D/2; i0 += warp_size) {
|
||||
const int i = i0 + threadIdx.x;
|
||||
if (i0 + warp_size > D/2 && i >= D/2) break;
|
||||
VKQ2[j*(D_padded/2) + i] *= KQ_max_scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j0 = 0; j0 < ncols; j0 += nwarps) {
|
||||
const int j_VKQ = j0 + threadIdx.y;
|
||||
if (ic0 + j_VKQ >= int(ne01.z)) {
|
||||
return;
|
||||
}
|
||||
|
||||
float KQ_rowsum_j;
|
||||
if (std::is_same<KQ_acc_t, float>::value) {
|
||||
KQ_rowsum_j = KQ_rowsum_f[j0/nwarps];
|
||||
} else {
|
||||
KQ_rowsum_j = __low2float(KQ_rowsum_h2[j0/nwarps]) + __high2float(KQ_rowsum_h2[j0/nwarps]);
|
||||
}
|
||||
|
||||
const int j_dst_unrolled = ((sequence*int(ne01.z) + ic0 + j_VKQ)*ne02 + head)*gridDim.y + blockIdx.y;
|
||||
|
||||
#pragma unroll
|
||||
for (int i0 = 0; i0 < D; i0 += warp_size) {
|
||||
const int i = i0 + threadIdx.x;
|
||||
if (i0 + warp_size > D && i >= D) {
|
||||
break;
|
||||
}
|
||||
float dst_val = VKQ[j_VKQ*D_padded + i];
|
||||
if (gridDim.y == 1) {
|
||||
dst_val /= KQ_rowsum_j;
|
||||
}
|
||||
dst[j_dst_unrolled*D + i] = dst_val;
|
||||
}
|
||||
|
||||
if (gridDim.y == 1 || threadIdx.x != 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
float2 dst_meta_val;
|
||||
if (std::is_same<KQ_acc_t, float>::value) {
|
||||
dst_meta_val.x = KQ_max_f[j0/nwarps];
|
||||
} else {
|
||||
dst_meta_val.x = __low2float(KQ_max_h2[j0/nwarps]);
|
||||
}
|
||||
dst_meta_val.y = KQ_rowsum_j;
|
||||
dst_meta[j_dst_unrolled] = dst_meta_val;
|
||||
}
|
||||
#else
|
||||
GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
|
||||
max_bias, m0, m1, n_head_log2, logit_softcap,
|
||||
ne00, ne01, ne02, ne03,
|
||||
nb01, nb02, nb03,
|
||||
ne10, ne11, ne12, ne13,
|
||||
nb11, nb12, nb13,
|
||||
nb21, nb22, nb23,
|
||||
ne31, ne32, ne33,
|
||||
nb31, nb32, nb33);
|
||||
NO_DEVICE_CODE;
|
||||
#endif // defined(FLASH_ATTN_AVAILABLE) && (defined(GGML_HIP_ROCWMMA_FATTN) && defined(GGML_USE_WMMA_FATTN))
|
||||
}
|
||||
|
||||
constexpr int get_max_power_of_2(int x) {
|
||||
return x % 2 == 0 ? 2*get_max_power_of_2(x/2) : 1;
|
||||
}
|
||||
|
||||
static_assert(get_max_power_of_2(1) == 1, "Test failed.");
|
||||
static_assert(get_max_power_of_2(2) == 2, "Test failed.");
|
||||
static_assert(get_max_power_of_2(4) == 4, "Test failed.");
|
||||
static_assert(get_max_power_of_2(6) == 2, "Test failed.");
|
||||
|
||||
// Number of VKQ rows calculated in parallel:
|
||||
constexpr int get_VKQ_stride(int D, int nwarps, int frag_m) {
|
||||
return (get_max_power_of_2(D/frag_m) < nwarps ? get_max_power_of_2(D/frag_m) : nwarps)*frag_m;
|
||||
}
|
||||
|
||||
static_assert(get_VKQ_stride(128, 1, 32) == 32, "Test failed.");
|
||||
static_assert(get_VKQ_stride(128, 2, 32) == 64, "Test failed.");
|
||||
static_assert(get_VKQ_stride(128, 4, 32) == 128, "Test failed.");
|
||||
static_assert(get_VKQ_stride( 64, 1, 32) == 32, "Test failed.");
|
||||
static_assert(get_VKQ_stride( 64, 2, 32) == 64, "Test failed.");
|
||||
static_assert(get_VKQ_stride( 64, 4, 32) == 64, "Test failed.");
|
||||
static_assert(get_VKQ_stride( 80, 1, 16) == 16, "Test failed.");
|
||||
static_assert(get_VKQ_stride( 80, 2, 16) == 16, "Test failed.");
|
||||
static_assert(get_VKQ_stride( 80, 4, 16) == 16, "Test failed.");
|
||||
|
||||
template <int D, int cols_per_block, typename KQ_acc_t>
|
||||
void ggml_cuda_flash_attn_ext_wmma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
||||
const ggml_tensor * KQV = dst;
|
||||
|
||||
constexpr int nwarps = 4;
|
||||
|
||||
constexpr int frag_m = cols_per_block == 8 && D % 32 == 0 ? 32 : 16;
|
||||
const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
|
||||
|
||||
float logit_softcap;
|
||||
memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));
|
||||
|
||||
fattn_kernel_t fattn_kernel;
|
||||
if (logit_softcap == 0.0f) {
|
||||
constexpr bool use_logit_softcap = false;
|
||||
fattn_kernel = flash_attn_ext_f16<
|
||||
D, cols_per_block, nwarps, get_VKQ_stride(D, nwarps, frag_m), KQ_acc_t, use_logit_softcap>;
|
||||
} else {
|
||||
constexpr bool use_logit_softcap = true;
|
||||
fattn_kernel = flash_attn_ext_f16<
|
||||
D, cols_per_block, nwarps, get_VKQ_stride(D, nwarps, frag_m), KQ_acc_t, use_logit_softcap>;
|
||||
}
|
||||
launch_fattn<D, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, 0, FATTN_KQ_STRIDE, true, true, false, warp_size);
|
||||
}
|
||||
|
||||
void ggml_cuda_flash_attn_ext_wmma_f16(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
||||
const ggml_tensor * KQV = dst;
|
||||
const ggml_tensor * Q = dst->src[0];
|
||||
|
||||
const enum ggml_prec prec = ggml_flash_attn_ext_get_prec(KQV);
|
||||
const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size;
|
||||
|
||||
if (prec != GGML_PREC_DEFAULT) {
|
||||
if (Q->ne[1] <= 32 || Q->ne[0] > 128) {
|
||||
constexpr int cols_per_block = 16;
|
||||
switch (Q->ne[0]) {
|
||||
case 64:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
case 80:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
case 96:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
case 112:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
case 128:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
case 256:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
default:
|
||||
GGML_ABORT("fatal error");
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
constexpr int cols_per_block = 32;
|
||||
switch (Q->ne[0]) {
|
||||
case 64:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
case 80:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
case 96:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
case 112:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
case 128:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, float>(ctx, dst);
|
||||
break;
|
||||
// case 256:
|
||||
// ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, float>(ctx, dst);
|
||||
// break;
|
||||
default:
|
||||
GGML_ABORT("fatal error");
|
||||
break;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
#if !defined(GGML_USE_HIP)
|
||||
if (Q->ne[1] <= 8 && Q->ne[0] % warp_size == 0) {
|
||||
constexpr int cols_per_block = 8;
|
||||
switch (Q->ne[0]) {
|
||||
case 64:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 96:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 128:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 256:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
default:
|
||||
GGML_ABORT("fatal error");
|
||||
break;
|
||||
}
|
||||
return;
|
||||
}
|
||||
#endif // !defined(GGML_USE_HIP)
|
||||
|
||||
if (Q->ne[1] <= 32) {
|
||||
constexpr int cols_per_block = 16;
|
||||
switch (Q->ne[0]) {
|
||||
case 64:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 80:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 96:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 112:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 128:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 256:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
default:
|
||||
GGML_ABORT("fatal error");
|
||||
break;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
constexpr int cols_per_block = 32;
|
||||
switch (Q->ne[0]) {
|
||||
case 64:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 64, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 80:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 80, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 96:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case< 96, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 112:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<112, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 128:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<128, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
case 256:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16_case<256, cols_per_block, half>(ctx, dst);
|
||||
break;
|
||||
default:
|
||||
GGML_ABORT("fatal error");
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -1,51 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
#include "common.cuh"
|
||||
|
||||
#if defined(GGML_USE_MUSA)
|
||||
#define GGML_USE_WMMA_FATTN
|
||||
#endif // defined(GGML_USE_MUSA)
|
||||
|
||||
#if defined(GGML_HIP_ROCWMMA_FATTN)
|
||||
#if defined(CDNA) && (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
|
||||
#define GGML_USE_WMMA_FATTN
|
||||
#elif defined(CDNA)
|
||||
#warning "rocwmma fattn on CDNA is broken on rocwmma v2.0.0, expect degraded performance"
|
||||
#endif // defined(CDNA) && (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
|
||||
#if defined(RDNA3)
|
||||
#define GGML_USE_WMMA_FATTN
|
||||
#endif // defined(RDNA3)
|
||||
#if defined(RDNA4) && ROCWMMA_VERSION_MAJOR > 1
|
||||
#define GGML_USE_WMMA_FATTN
|
||||
#elif defined(RDNA4)
|
||||
#warning "rocwmma fattn is not supported on RDNA4 on rocwmma < v2.0.0, expect degraded performance"
|
||||
#endif // defined(RDNA4) && ROCWMMA_VERSION_MAJOR > 1
|
||||
#endif // defined(GGML_HIP_ROCWMMA_FATTN)
|
||||
|
||||
// WMMA flash attention requires FP16 matrix instructions to be available for ggml code.
|
||||
static bool ggml_cuda_should_use_wmma_fattn(const int cc) {
|
||||
#if defined(GGML_USE_HIP) && !defined(GGML_HIP_ROCWMMA_FATTN)
|
||||
return false;
|
||||
#else
|
||||
if ((GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_VOLTA) ||
|
||||
GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_MTHREADS(cc)) {
|
||||
return true;
|
||||
} else if (GGML_CUDA_CC_IS_CDNA(cc)){
|
||||
#if defined(GGML_HIP_ROCWMMA_FATTN) && (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
#endif // defined(GGML_HIP_ROCWMMA_FATTN) (ROCWMMA_VERSION_MAJOR < 2 || ROCWMMA_VERSION_MINOR > 0 || ROCWMMA_VERSION_PATCH > 0)
|
||||
} else if (GGML_CUDA_CC_IS_RDNA4(cc)) {
|
||||
#if defined(GGML_HIP_ROCWMMA_FATTN) && ROCWMMA_VERSION_MAJOR > 1
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
#endif // defined(GGML_HIP_ROCWMMA_FATTN) && ROCWMMA_VERSION_MAJOR > 1
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
#endif // defined(GGML_USE_HIP) && !defined(GGML_HIP_ROCWMMA_FATTN)
|
||||
}
|
||||
|
||||
void ggml_cuda_flash_attn_ext_wmma_f16(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
|
||||
@@ -3,7 +3,6 @@
|
||||
#include "fattn-mma-f16.cuh"
|
||||
#include "fattn-tile.cuh"
|
||||
#include "fattn-vec.cuh"
|
||||
#include "fattn-wmma-f16.cuh"
|
||||
#include "fattn.cuh"
|
||||
|
||||
template <int DKQ, int DV, int ncols2>
|
||||
@@ -330,11 +329,10 @@ static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_t
|
||||
|
||||
// Best FlashAttention kernel for a specific GPU:
|
||||
enum best_fattn_kernel {
|
||||
BEST_FATTN_KERNEL_NONE = 0,
|
||||
BEST_FATTN_KERNEL_TILE = 200,
|
||||
BEST_FATTN_KERNEL_VEC = 100,
|
||||
BEST_FATTN_KERNEL_WMMA_F16 = 300,
|
||||
BEST_FATTN_KERNEL_MMA_F16 = 400,
|
||||
BEST_FATTN_KERNEL_NONE = 0,
|
||||
BEST_FATTN_KERNEL_TILE = 200,
|
||||
BEST_FATTN_KERNEL_VEC = 100,
|
||||
BEST_FATTN_KERNEL_MMA_F16 = 400,
|
||||
};
|
||||
|
||||
static bool ggml_cuda_fattn_kv_type_supported(ggml_type type) {
|
||||
@@ -500,14 +498,6 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
|
||||
return BEST_FATTN_KERNEL_MMA_F16;
|
||||
}
|
||||
|
||||
// Use the WMMA kernel if possible:
|
||||
if (ggml_cuda_should_use_wmma_fattn(cc) && K->ne[1] % FATTN_KQ_STRIDE == 0 && Q->ne[0] != 40 && Q->ne[0] != 72 && Q->ne[0] != 192 && Q->ne[0] != 512 && Q->ne[0] != 576) {
|
||||
if (can_use_vector_kernel && Q->ne[1] <= 2) {
|
||||
return BEST_FATTN_KERNEL_VEC;
|
||||
}
|
||||
return BEST_FATTN_KERNEL_WMMA_F16;
|
||||
}
|
||||
|
||||
// AMD MFMA needs a certain minimum batch size to outscale the tile kernel for large head sizes.
|
||||
if ((amd_mfma_available(cc) && Q->ne[0] <= 256) && Q->ne[0] != 40 && Q->ne[0] != 72) {
|
||||
if ((Q->ne[0] <= 64 && Q->ne[1] * gqa_ratio_eff > 8)) {
|
||||
@@ -559,7 +549,6 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d
|
||||
|
||||
switch (kernel) {
|
||||
case BEST_FATTN_KERNEL_TILE:
|
||||
case BEST_FATTN_KERNEL_WMMA_F16:
|
||||
case BEST_FATTN_KERNEL_MMA_F16:
|
||||
need_f16_K = true;
|
||||
need_f16_V = true;
|
||||
@@ -589,9 +578,6 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
|
||||
case BEST_FATTN_KERNEL_VEC:
|
||||
ggml_cuda_flash_attn_ext_vec(ctx, dst);
|
||||
break;
|
||||
case BEST_FATTN_KERNEL_WMMA_F16:
|
||||
ggml_cuda_flash_attn_ext_wmma_f16(ctx, dst);
|
||||
break;
|
||||
case BEST_FATTN_KERNEL_MMA_F16:
|
||||
ggml_cuda_flash_attn_ext_mma_f16(ctx, dst);
|
||||
break;
|
||||
|
||||
@@ -9,6 +9,21 @@ using namespace cub;
|
||||
|
||||
#include "ssm-scan.cuh"
|
||||
|
||||
|
||||
// Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan path.
|
||||
// For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences).
|
||||
#define SSM_SSD_MIN_TOKENS 128
|
||||
|
||||
// prepare_dt kernel dimensions: one block per (head, seq), each block handles DT_MAX_ITEMS items.
|
||||
#define SSM_SSD_DT_BLOCK 256
|
||||
#define SSM_SSD_DT_MAX_ITEMS 32
|
||||
|
||||
// Maximum tokens the SSD path supports, derived from the prepare_dt kernel block capacity.
|
||||
#define SSM_SSD_MAX_TOKENS (SSM_SSD_DT_BLOCK * SSM_SSD_DT_MAX_ITEMS)
|
||||
|
||||
// Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk.
|
||||
#define SSM_SSD_CHUNK_SIZE 256
|
||||
|
||||
// We would like to keep pragma unroll for cases where L_template is not 0,
|
||||
// so we suppress the clang transformation warning.
|
||||
#ifdef __clang__
|
||||
@@ -316,6 +331,429 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
|
||||
}
|
||||
}
|
||||
|
||||
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
|
||||
// ============================================================================
|
||||
// SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS)
|
||||
//
|
||||
// Instead of a sequential scan, SSD reformulates the output as:
|
||||
// Y = (L (.) (C @ B^T)) @ (X * dt) + decay * C @ s_init
|
||||
// where L is a causal decay mask derived from A and dt.
|
||||
//
|
||||
// This converts the O(T*N) sequential scan into parallel matmuls.
|
||||
// ============================================================================
|
||||
// Softplus(dt) and inclusive prefix sum per head using CUB BlockScan.
|
||||
// Grid: (n_head, n_seqs)
|
||||
template <int BLOCK_SIZE, int MAX_ITEMS>
|
||||
__global__ void ssm_ssd_prepare_dt_kernel(
|
||||
const float * __restrict__ dt_raw,
|
||||
float * __restrict__ dt_sp_out,
|
||||
float * __restrict__ cs_out,
|
||||
const int n_head, const int n_tok,
|
||||
const int dt_stride_tok, // elements between tokens in dt
|
||||
const int dt_stride_seq) { // elements between sequences in dt
|
||||
|
||||
const int h = blockIdx.x;
|
||||
const int s = blockIdx.y;
|
||||
|
||||
const float * dt_seq = dt_raw + s * dt_stride_seq;
|
||||
|
||||
float * dt_sp_seq = dt_sp_out + s * n_tok * n_head;
|
||||
float * cs_seq = cs_out + s * n_tok * n_head;
|
||||
|
||||
const int items_per_thread = (n_tok + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
||||
|
||||
// Phase 1: softplus with interleaved distribution (t = i*BLOCK_SIZE + threadIdx.x).
|
||||
// Each warp reads BLOCK_SIZE consecutive tokens, giving coalesced dt_raw loads
|
||||
// (stride n_head between threads vs. items_per_thread*n_head in blocked layout).
|
||||
float local_vals[MAX_ITEMS];
|
||||
for (int i = 0; i < items_per_thread; i++) {
|
||||
const int t = i * BLOCK_SIZE + threadIdx.x;
|
||||
if (t < n_tok) {
|
||||
float val = dt_seq[h + t * dt_stride_tok];
|
||||
float sp = (val <= 20.0f) ? log1pf(expf(val)) : val;
|
||||
local_vals[i] = sp;
|
||||
dt_sp_seq[t * n_head + h] = sp;
|
||||
} else {
|
||||
local_vals[i] = 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
// Phase 2+3: per-step inclusive scan to build cs[] in token order.
|
||||
// With interleaved distribution the per-thread total scan would not give token-order
|
||||
// prefix sums, so we scan one BLOCK_SIZE slab at a time and carry a running total.
|
||||
#ifdef USE_CUB
|
||||
using BlockScan = cub::BlockScan<float, BLOCK_SIZE>;
|
||||
__shared__ typename BlockScan::TempStorage scan_temp;
|
||||
__shared__ float step_total;
|
||||
|
||||
float running = 0.0f;
|
||||
for (int i = 0; i < items_per_thread; i++) {
|
||||
float inclusive;
|
||||
BlockScan(scan_temp).InclusiveSum(local_vals[i], inclusive);
|
||||
const int t = i * BLOCK_SIZE + threadIdx.x;
|
||||
if (t < n_tok) {
|
||||
cs_seq[t * n_head + h] = running + inclusive;
|
||||
}
|
||||
if (threadIdx.x == BLOCK_SIZE - 1) {
|
||||
step_total = inclusive;
|
||||
}
|
||||
__syncthreads();
|
||||
running += step_total;
|
||||
}
|
||||
#else
|
||||
// Fallback: sequential prefix scan in shared memory, one slab at a time.
|
||||
__shared__ float sdata[BLOCK_SIZE];
|
||||
float running = 0.0f;
|
||||
for (int i = 0; i < items_per_thread; i++) {
|
||||
const int t = i * BLOCK_SIZE + threadIdx.x;
|
||||
sdata[threadIdx.x] = local_vals[i];
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0) {
|
||||
for (int j = 1; j < BLOCK_SIZE; j++) {
|
||||
sdata[j] += sdata[j - 1];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
if (t < n_tok) {
|
||||
cs_seq[t * n_head + h] = running + sdata[threadIdx.x];
|
||||
}
|
||||
running += sdata[BLOCK_SIZE - 1];
|
||||
__syncthreads();
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// Prepare SSD matmul inputs for one chunk: X_dt, B_weighted, C_scaled.
|
||||
// T_matmul controls precision for X_dt, B_weighted (float or half).
|
||||
// C_scaled is always float (pairs with float s_cur in step 3c).
|
||||
// Computation is always FP32; only the final store converts to T_matmul.
|
||||
// Also materializes the causal M matrix = exp(A*(cs_out - cs_in)) * CB (fused with prep to save a launch).
|
||||
// Grid: (ceil(max(C*head_dim, d_state*C, chunk_len^2) / BLOCK), n_head, n_seqs)
|
||||
template <int BLOCK_SIZE, typename T_matmul>
|
||||
__global__ void ssm_ssd_pre_matmul_kernel(
|
||||
const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums
|
||||
const float * __restrict__ dt_sp, // {n_tok, n_head} softplus(dt)
|
||||
const float * __restrict__ A, // {1, n_head}
|
||||
const float * __restrict__ x, // {head_dim, n_head, n_tok, n_seqs}
|
||||
const float * __restrict__ B, // {d_state, n_group, n_tok, n_seqs}
|
||||
const float * __restrict__ C_src, // {d_state, n_group, n_tok, n_seqs}
|
||||
T_matmul * __restrict__ X_dt, // {head_dim, C, n_head} x * dt, d-fastest
|
||||
T_matmul * __restrict__ B_weighted, // {d_state, C, n_head} B * decay_from_end
|
||||
float * __restrict__ C_scaled, // {d_state, C, n_head} C * decay_to_pos (always float)
|
||||
const float * __restrict__ CB, // {chunk_len, chunk_len, n_group, n_seqs}
|
||||
half * __restrict__ M_out, // {chunk_len, chunk_len, n_head, n_seqs}
|
||||
const int chunk_len, const int head_dim, const int n_head, const int n_group,
|
||||
const int d_state, const int A_stride,
|
||||
const int x_stride_tok, const int x_stride_seq,
|
||||
const int B_stride_tok, const int B_stride_seq,
|
||||
const int C_stride_tok, const int C_stride_seq,
|
||||
const int chunk_offset,
|
||||
const int n_tok_total) {
|
||||
|
||||
const int h = blockIdx.y;
|
||||
const int s = blockIdx.z;
|
||||
const int g = h / (n_head / n_group);
|
||||
|
||||
const float A_h = A[h * A_stride];
|
||||
const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x;
|
||||
|
||||
const int cs_seq_off = s * n_tok_total * n_head;
|
||||
const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f;
|
||||
const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base;
|
||||
|
||||
// Prepare X_dt = x * dt, stored d-fastest for coalesced reads and writes.
|
||||
const int n_xdt = chunk_len * head_dim;
|
||||
if (idx < n_xdt) {
|
||||
const int d = idx % head_dim;
|
||||
const int t = idx / head_dim;
|
||||
|
||||
const float x_val = x[s * x_stride_seq + (chunk_offset + t) * x_stride_tok + d + h * head_dim];
|
||||
const float dt_val = dt_sp[cs_seq_off + (chunk_offset + t) * n_head + h];
|
||||
|
||||
X_dt[d + t * head_dim + h * n_xdt + s * n_xdt * n_head] = (T_matmul)(x_val * dt_val);
|
||||
}
|
||||
|
||||
// Prepare B_weighted and C_scaled together: both share the same index space (d_state * chunk_len)
|
||||
// and the same cs_t load, so merging halves the cs[] global memory traffic.
|
||||
const int n_bw = d_state * chunk_len;
|
||||
if (idx < n_bw) {
|
||||
const int n = idx % d_state;
|
||||
const int t = idx / d_state;
|
||||
|
||||
const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base;
|
||||
|
||||
const float B_val = B[s * B_stride_seq + (chunk_offset + t) * B_stride_tok + g * d_state + n];
|
||||
B_weighted[n + t * d_state + h * n_bw + s * n_bw * n_head] = (T_matmul)(B_val * __expf(A_h * (cs_last - cs_t)));
|
||||
|
||||
const float C_val = C_src[s * C_stride_seq + (chunk_offset + t) * C_stride_tok + g * d_state + n];
|
||||
C_scaled[n + t * d_state + h * n_bw + s * n_bw * n_head] = C_val * __expf(A_h * cs_t);
|
||||
}
|
||||
|
||||
// Materialize M = exp(A*(cs_out - cs_in)) * CB with causal mask.
|
||||
const int n_M = chunk_len * chunk_len;
|
||||
if (idx < n_M) {
|
||||
const int t_out = idx % chunk_len;
|
||||
const int t_in = idx / chunk_len;
|
||||
|
||||
half val;
|
||||
if (t_in <= t_out) {
|
||||
const float cs_out = cs[cs_seq_off + (chunk_offset + t_out) * n_head + h] - cs_base;
|
||||
const float cs_in = cs[cs_seq_off + (chunk_offset + t_in) * n_head + h] - cs_base;
|
||||
const float decay = __expf(A_h * (cs_out - cs_in));
|
||||
const float * CB_g = CB + (int64_t)s * chunk_len * chunk_len * n_group
|
||||
+ (int64_t)g * chunk_len * chunk_len;
|
||||
const float cb_val = CB_g[t_out + t_in * chunk_len];
|
||||
val = __float2half(decay * cb_val);
|
||||
} else {
|
||||
val = __float2half(0.0f);
|
||||
}
|
||||
|
||||
M_out[(int64_t)s * n_M * n_head + (int64_t)h * n_M + t_in * chunk_len + t_out] = val;
|
||||
}
|
||||
}
|
||||
|
||||
// Scale running state in-place: s_cur *= decay_total(chunk).
|
||||
// Called BEFORE cuBLAS state update (beta=1) to fuse inter-chunk decay.
|
||||
// Eliminates the s_old buffer and D2D memcpy vs the old approach of:
|
||||
// memcpy(s_old, s_cur) -> cuBLAS(beta=0) -> s_cur += decay * s_old
|
||||
// Grid: (ceil(d_state * head_dim / BLOCK), n_head, n_seqs)
|
||||
template <int BLOCK_SIZE>
|
||||
__global__ void ssm_ssd_scale_state_kernel(
|
||||
float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs}
|
||||
const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums
|
||||
const float * __restrict__ A, // {1, n_head}
|
||||
const int d_state, const int head_dim, const int n_head,
|
||||
const int chunk_offset, const int chunk_len,
|
||||
const int n_tok_total, const int A_stride) {
|
||||
|
||||
const int h = blockIdx.y;
|
||||
const int s = blockIdx.z;
|
||||
const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x;
|
||||
const int state_per_head = d_state * head_dim;
|
||||
if (idx >= state_per_head) return;
|
||||
|
||||
const float A_h = A[h * A_stride];
|
||||
const int cs_seq_off = s * n_tok_total * n_head;
|
||||
const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f;
|
||||
const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base;
|
||||
const float decay_total = __expf(A_h * cs_last);
|
||||
|
||||
const int off = s * state_per_head * n_head + h * state_per_head + idx;
|
||||
s_cur[off] *= decay_total;
|
||||
}
|
||||
|
||||
// Copy initial state from src0[ids[s]] into s_cur for each sequence.
|
||||
// Grid: (ceil(d_state * head_dim * n_head / BLOCK), n_seqs)
|
||||
template <int BLOCK_SIZE>
|
||||
__global__ void ssm_ssd_init_state_kernel(
|
||||
const float * __restrict__ src0, // {d_state, head_dim, n_head, n_rs}
|
||||
const int32_t * __restrict__ ids, // {n_seqs}
|
||||
float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs}
|
||||
const int state_size, // d_state * head_dim * n_head
|
||||
const int64_t s0_stride_seq) { // elements between state rows
|
||||
const int s = blockIdx.y;
|
||||
const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x;
|
||||
if (idx >= state_size) return;
|
||||
|
||||
const float * s_src = src0 + (int64_t)ids[s] * s0_stride_seq;
|
||||
s_cur[s * state_size + idx] = s_src[idx];
|
||||
}
|
||||
|
||||
// SSD (State Space Duality) dispatch for Mamba-2 prefill.
|
||||
// Chunked matmuls: CB, materialize M + cuBLAS Y, S@C, B@X_dt.
|
||||
// All strides are in elements (floats), not bytes.
|
||||
static void ssm_scan_ssd_f32_cuda(
|
||||
ggml_backend_cuda_context & ctx,
|
||||
const float * src0_d, const float * src1_d, const float * src2_d, const float * src3_d,
|
||||
const float * src4_d, const float * src5_d, const int32_t * src6_d, float * dst_d,
|
||||
const int64_t s0_stride_seq, // state (src0) stride between seqs
|
||||
const int x_stride_tok, const int x_stride_seq, // x (src1) strides
|
||||
const int dt_stride_tok, const int dt_stride_seq, // dt (src2) strides
|
||||
const int A_stride, // A (src3) stride between heads
|
||||
const int B_stride_tok, const int B_stride_seq, // B (src4) strides
|
||||
const int C_stride_tok, const int C_stride_seq, // C (src5) strides
|
||||
const int64_t s_off, const int64_t d_state, const int64_t head_dim,
|
||||
const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq) {
|
||||
|
||||
cudaStream_t stream = ctx.stream();
|
||||
const int64_t d_inner = head_dim * n_head;
|
||||
|
||||
const int64_t chunk_size = SSM_SSD_CHUNK_SIZE;
|
||||
const int64_t n_chunks = (n_tok + chunk_size - 1) / chunk_size;
|
||||
|
||||
const int64_t state_per_head = d_state * head_dim;
|
||||
|
||||
using matmul_t = half;
|
||||
static constexpr cudaDataType_t matmul_dtype = CUDA_R_16F;
|
||||
|
||||
ggml_cuda_pool_alloc<float> dt_sp_buf(ctx.pool(), n_tok * n_head * n_seq);
|
||||
ggml_cuda_pool_alloc<float> cs_buf(ctx.pool(), n_tok * n_head * n_seq);
|
||||
ggml_cuda_pool_alloc<float> CB_buf(ctx.pool(), chunk_size * chunk_size * n_group * n_seq);
|
||||
ggml_cuda_pool_alloc<matmul_t> X_dt_buf(ctx.pool(), chunk_size * head_dim * n_head * n_seq);
|
||||
ggml_cuda_pool_alloc<matmul_t> B_w_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq);
|
||||
ggml_cuda_pool_alloc<float> C_s_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq);
|
||||
float * dt_sp = dt_sp_buf.get();
|
||||
float * cs = cs_buf.get();
|
||||
float * CB = CB_buf.get();
|
||||
matmul_t * X_dt = X_dt_buf.get();
|
||||
matmul_t * B_weighted = B_w_buf.get();
|
||||
float * C_scaled = C_s_buf.get();
|
||||
float * s_cur = (float *)((char *)dst_d + s_off); // write state directly to dst
|
||||
|
||||
// Step 1: softplus(dt) and parallel prefix sum over full sequence
|
||||
{
|
||||
dim3 grid(n_head, n_seq);
|
||||
ssm_ssd_prepare_dt_kernel<SSM_SSD_DT_BLOCK, SSM_SSD_DT_MAX_ITEMS><<<grid, SSM_SSD_DT_BLOCK, 0, stream>>>(
|
||||
src2_d, dt_sp, cs, n_head, n_tok, dt_stride_tok, dt_stride_seq);
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
}
|
||||
|
||||
// Step 2: initialize running state from src0[ids[s]]
|
||||
{
|
||||
constexpr int BLOCK = 256;
|
||||
const int64_t state_size = d_state * head_dim * n_head;
|
||||
dim3 grid((state_size + BLOCK - 1) / BLOCK, n_seq);
|
||||
ssm_ssd_init_state_kernel<BLOCK><<<grid, BLOCK, 0, stream>>>(
|
||||
src0_d, src6_d, s_cur, state_size, s0_stride_seq);
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
}
|
||||
|
||||
// Step 3: chunked SSD loop
|
||||
// Per chunk: pre_matmul (incl. M) + 4 cuBLAS (CB, Y, S@C, state update) + scale_state
|
||||
cublasHandle_t handle = ctx.cublas_handle();
|
||||
CUBLAS_CHECK(cublasSetStream(handle, stream));
|
||||
const float alpha_one = 1.0f;
|
||||
const float beta_zero = 0.0f;
|
||||
const float beta_one = 1.0f;
|
||||
const int lda_C_src = C_stride_tok; // leading dim for C in CB = C^T @ B
|
||||
const int ldb_B_src = B_stride_tok; // leading dim for B in CB = C^T @ B
|
||||
|
||||
// Scratch buffer for causal M matrix, reused across chunks (max size at chunk_size)
|
||||
const int64_t n_M_max = chunk_size * chunk_size;
|
||||
ggml_cuda_pool_alloc<half> M_buf(ctx.pool(), n_M_max * n_head * n_seq);
|
||||
half * M_mat = M_buf.get();
|
||||
|
||||
for (int64_t k = 0; k < n_chunks; k++) {
|
||||
const int64_t chunk_offset = k * chunk_size;
|
||||
const int64_t chunk_len = (chunk_offset + chunk_size <= n_tok) ? chunk_size : (n_tok - chunk_offset);
|
||||
|
||||
// 3a: CB = C^T @ B per group
|
||||
for (int64_t s = 0; s < n_seq; s++) {
|
||||
const float * C_s = src5_d + s * C_stride_seq + chunk_offset * C_stride_tok;
|
||||
const float * B_s = src4_d + s * B_stride_seq + chunk_offset * B_stride_tok;
|
||||
float * CB_s = CB + s * chunk_len * chunk_len * n_group;
|
||||
|
||||
if (n_group == 1) {
|
||||
CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
chunk_len, chunk_len, d_state,
|
||||
&alpha_one, C_s, lda_C_src, B_s, ldb_B_src,
|
||||
&beta_zero, CB_s, (int)chunk_len));
|
||||
} else {
|
||||
CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
chunk_len, chunk_len, d_state,
|
||||
&alpha_one,
|
||||
C_s, CUDA_R_32F, lda_C_src, d_state,
|
||||
B_s, CUDA_R_32F, ldb_B_src, d_state,
|
||||
&beta_zero,
|
||||
CB_s, CUDA_R_32F, (int)chunk_len, (long long)(chunk_len * chunk_len),
|
||||
n_group,
|
||||
CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
|
||||
}
|
||||
}
|
||||
|
||||
// 3b: prepare X_dt, B_weighted, C_scaled + materialize causal M matrix
|
||||
const int64_t n_M = chunk_len * chunk_len;
|
||||
{
|
||||
constexpr int BLOCK = 256;
|
||||
const int64_t n_xdt = chunk_len * head_dim;
|
||||
const int64_t n_bw = d_state * chunk_len;
|
||||
int64_t max_work = n_xdt;
|
||||
if (n_bw > max_work) max_work = n_bw;
|
||||
if (n_M > max_work) max_work = n_M;
|
||||
dim3 grid((max_work + BLOCK - 1) / BLOCK, n_head, n_seq);
|
||||
ssm_ssd_pre_matmul_kernel<BLOCK, matmul_t><<<grid, BLOCK, 0, stream>>>(
|
||||
cs, dt_sp, src3_d, src1_d, src4_d, src5_d,
|
||||
X_dt, B_weighted, C_scaled,
|
||||
CB, M_mat,
|
||||
chunk_len, head_dim, n_head, n_group, d_state, A_stride,
|
||||
x_stride_tok, x_stride_seq, B_stride_tok, B_stride_seq, C_stride_tok, C_stride_seq,
|
||||
chunk_offset, n_tok);
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
}
|
||||
|
||||
// 3c: dst = S_cur^T @ C_scaled (state contribution)
|
||||
{
|
||||
const int64_t stride_S = state_per_head;
|
||||
const int64_t stride_Cs = d_state * chunk_len;
|
||||
|
||||
for (int64_t s = 0; s < n_seq; s++) {
|
||||
float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner;
|
||||
|
||||
CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
head_dim, chunk_len, d_state,
|
||||
&alpha_one,
|
||||
s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S,
|
||||
C_scaled + s * stride_Cs * n_head, CUDA_R_32F, d_state, stride_Cs,
|
||||
&beta_zero,
|
||||
dst_chunk, CUDA_R_32F, d_inner, head_dim,
|
||||
n_head,
|
||||
CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
|
||||
}
|
||||
}
|
||||
|
||||
// 3d: dst += X_dt @ M^T (intra-chunk contribution, adds to 3c result)
|
||||
// M is stored as M[t_out, t_in] (lower-triangular), transpose needed for Y = X @ M^T.
|
||||
{
|
||||
const int64_t stride_M = n_M;
|
||||
const int64_t stride_X_h = (int64_t)chunk_len * head_dim;
|
||||
|
||||
for (int64_t s = 0; s < n_seq; s++) {
|
||||
float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner;
|
||||
CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T,
|
||||
head_dim, chunk_len, chunk_len,
|
||||
&alpha_one,
|
||||
X_dt + s * stride_X_h * n_head, matmul_dtype, head_dim, stride_X_h,
|
||||
M_mat + s * stride_M * n_head, matmul_dtype, chunk_len, stride_M,
|
||||
&beta_one,
|
||||
dst_chunk, CUDA_R_32F, d_inner, head_dim,
|
||||
n_head,
|
||||
CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
|
||||
}
|
||||
}
|
||||
|
||||
// 3e: s_cur = B_weighted @ X_dt^T + decay_total * s_cur_old (state update)
|
||||
{
|
||||
// Scale s_cur in-place by per-head decay_total BEFORE cuBLAS overwrites it
|
||||
constexpr int BLOCK = 256;
|
||||
dim3 grid((state_per_head + BLOCK - 1) / BLOCK, n_head, n_seq);
|
||||
ssm_ssd_scale_state_kernel<BLOCK><<<grid, BLOCK, 0, stream>>>(
|
||||
s_cur, cs, src3_d,
|
||||
d_state, head_dim, n_head,
|
||||
chunk_offset, chunk_len, n_tok, A_stride);
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
|
||||
// cuBLAS with beta=1: s_cur = B_weighted @ X_dt^T + 1.0 * s_cur (already scaled)
|
||||
const int64_t stride_Bw = d_state * chunk_len;
|
||||
const int64_t stride_X = chunk_len * head_dim;
|
||||
const int64_t stride_S = state_per_head;
|
||||
|
||||
for (int64_t s = 0; s < n_seq; s++) {
|
||||
// X_dt is d-fastest {hd, C}, read as OP_T to get {C, hd}
|
||||
CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T,
|
||||
d_state, head_dim, chunk_len,
|
||||
&alpha_one,
|
||||
B_weighted + s * stride_Bw * n_head, matmul_dtype, d_state, stride_Bw,
|
||||
X_dt + s * stride_X * n_head, matmul_dtype, head_dim, stride_X,
|
||||
&beta_one,
|
||||
s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S,
|
||||
n_head,
|
||||
CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
|
||||
|
||||
void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
||||
const struct ggml_tensor * src0 = dst->src[0]; // s
|
||||
const struct ggml_tensor * src1 = dst->src[1]; // x
|
||||
@@ -357,6 +795,49 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
||||
GGML_ASSERT(src6->type == GGML_TYPE_I32);
|
||||
GGML_ASSERT(dst->type == GGML_TYPE_F32);
|
||||
|
||||
// Byte strides are narrowed to int for both scan and SSD paths.
|
||||
GGML_ASSERT(src0->nb[2] <= (size_t)INT_MAX);
|
||||
GGML_ASSERT(src0->nb[3] <= (size_t)INT_MAX);
|
||||
GGML_ASSERT(src1->nb[2] <= (size_t)INT_MAX);
|
||||
GGML_ASSERT(src1->nb[3] <= (size_t)INT_MAX);
|
||||
GGML_ASSERT(src2->nb[1] <= (size_t)INT_MAX);
|
||||
GGML_ASSERT(src2->nb[2] <= (size_t)INT_MAX);
|
||||
GGML_ASSERT(src3->nb[1] <= (size_t)INT_MAX);
|
||||
GGML_ASSERT(src4->nb[2] <= (size_t)INT_MAX);
|
||||
GGML_ASSERT(src4->nb[3] <= (size_t)INT_MAX);
|
||||
GGML_ASSERT(src5->nb[2] <= (size_t)INT_MAX);
|
||||
GGML_ASSERT(src5->nb[3] <= (size_t)INT_MAX);
|
||||
|
||||
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
|
||||
// Mamba-2 with scalar A per head: use SSD matmul path for long sequences.
|
||||
// Requires NVIDIA Turing+ otherwise fallback to scan.
|
||||
const bool is_mamba2 = (src3->nb[1] == sizeof(float));
|
||||
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
|
||||
const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS
|
||||
&& n_t <= SSM_SSD_MAX_TOKENS
|
||||
&& GGML_CUDA_CC_IS_NVIDIA(cc)
|
||||
&& cc >= GGML_CUDA_CC_TURING
|
||||
&& nr % 8 == 0; // cuBLAS requires 8-element (16-byte) alignment
|
||||
|
||||
if (use_ssd) {
|
||||
// ssm_ssd_init_state_kernel uses flat linear indexing within each sequence,
|
||||
// so src0 must be fully contiguous across all inner dimensions.
|
||||
// The scan path handles non-contiguous nb[2] via src0_nb2 but does not handle nb[1].
|
||||
GGML_ASSERT(src0->nb[1] == nc * sizeof(float));
|
||||
GGML_ASSERT(src0->nb[2] == nc * nr * sizeof(float));
|
||||
|
||||
ssm_scan_ssd_f32_cuda(ctx,
|
||||
src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d,
|
||||
(int64_t)(src0->nb[3] / sizeof(float)),
|
||||
(int)(src1->nb[2] / sizeof(float)), (int)(src1->nb[3] / sizeof(float)),
|
||||
(int)(src2->nb[1] / sizeof(float)), (int)(src2->nb[2] / sizeof(float)),
|
||||
(int)(src3->nb[1] / sizeof(float)),
|
||||
(int)(src4->nb[2] / sizeof(float)), (int)(src4->nb[3] / sizeof(float)),
|
||||
(int)(src5->nb[2] / sizeof(float)), (int)(src5->nb[3] / sizeof(float)),
|
||||
s_off, nc, nr, nh, ng, n_t, n_s);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d,
|
||||
src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2],
|
||||
src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3],
|
||||
|
||||
Vendored
-4
@@ -6,10 +6,6 @@
|
||||
#include <hip/hip_fp16.h>
|
||||
#include <hip/hip_bf16.h>
|
||||
|
||||
#if defined(GGML_HIP_ROCWMMA_FATTN)
|
||||
#include <rocwmma/rocwmma-version.hpp>
|
||||
#endif // defined(GGML_HIP_ROCWMMA_FATTN)
|
||||
|
||||
#ifdef GGML_USE_NCCL
|
||||
#include <rccl/rccl.h>
|
||||
#endif // GGML_USE_NCCL
|
||||
|
||||
@@ -3281,6 +3281,35 @@ static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session *
|
||||
GGML_UNUSED(sess);
|
||||
}
|
||||
|
||||
static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
|
||||
const struct ggml_tensor * src1 = op->src[1];
|
||||
const struct ggml_tensor * dst = op;
|
||||
|
||||
const bool is_2D = ((const int32_t *) op->op_params)[6] == 1;
|
||||
if (!is_2D) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// For now support F32->F32 and F32->F16 only.
|
||||
if (src1->type != GGML_TYPE_F32 || (dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_F32)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// For now keep padded OPs on CPU. Will revisit once we expand coverage past patch-embed OPs.
|
||||
const int32_t p0 = ((const int32_t *) op->op_params)[2];
|
||||
const int32_t p1 = ((const int32_t *) op->op_params)[3];
|
||||
if (p0 != 0 || p1 != 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
GGML_UNUSED(sess);
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool ggml_hexagon_supported_pad(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
|
||||
const struct ggml_tensor * src0 = op->src[0];
|
||||
const struct ggml_tensor * dst = op;
|
||||
@@ -3430,6 +3459,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
|
||||
case GGML_OP_SOLVE_TRI: return HTP_OP_SOLVE_TRI;
|
||||
case GGML_OP_TRI: return HTP_OP_TRI;
|
||||
case GGML_OP_PAD: return HTP_OP_PAD;
|
||||
case GGML_OP_IM2COL: return HTP_OP_IM2COL;
|
||||
|
||||
case GGML_OP_UNARY:
|
||||
switch (ggml_get_unary_op(t)) {
|
||||
@@ -4152,6 +4182,10 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
|
||||
supp = ggml_hexagon_supported_ssm_conv(sess, op);
|
||||
break;
|
||||
|
||||
case GGML_OP_IM2COL:
|
||||
supp = ggml_hexagon_supported_im2col(sess, op);
|
||||
break;
|
||||
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
supp = ggml_hexagon_supported_gated_delta_net(sess, op);
|
||||
break;
|
||||
|
||||
@@ -42,6 +42,7 @@ add_library(${HTP_LIB} SHARED
|
||||
solve-tri-ops.c
|
||||
pad-ops.c
|
||||
argsort-ops.c
|
||||
im2col-ops.c
|
||||
)
|
||||
|
||||
target_compile_definitions(${HTP_LIB} PRIVATE
|
||||
|
||||
@@ -140,5 +140,6 @@ int op_diag(struct htp_ops_context * octx);
|
||||
int op_solve_tri(struct htp_ops_context * octx);
|
||||
int op_gated_delta_net(struct htp_ops_context * octx);
|
||||
int op_pad(struct htp_ops_context * octx);
|
||||
int op_im2col(struct htp_ops_context * octx);
|
||||
|
||||
#endif /* HTP_CTX_H */
|
||||
|
||||
@@ -98,6 +98,7 @@ enum htp_op_code {
|
||||
HTP_OP_NORM,
|
||||
HTP_OP_CONCAT,
|
||||
HTP_OP_CLAMP,
|
||||
HTP_OP_IM2COL,
|
||||
|
||||
HTP_OP_INVALID
|
||||
};
|
||||
|
||||
@@ -0,0 +1,306 @@
|
||||
#pragma clang diagnostic ignored "-Wunused-variable"
|
||||
#pragma clang diagnostic ignored "-Wunused-function"
|
||||
#pragma clang diagnostic ignored "-Wunused-but-set-variable"
|
||||
|
||||
#include <HAP_farf.h>
|
||||
#include <HAP_perf.h>
|
||||
#include <hexagon_protos.h>
|
||||
#include <hexagon_types.h>
|
||||
#include <string.h>
|
||||
|
||||
#define GGML_COMMON_DECL_C
|
||||
#include "ggml-common.h"
|
||||
#include "htp-ctx.h"
|
||||
#include "htp-ops.h"
|
||||
#include "hvx-utils.h"
|
||||
#include "hex-dma.h"
|
||||
#include "hex-profile.h"
|
||||
#include "htp-vtcm.h"
|
||||
|
||||
struct htp_im2col_context {
|
||||
struct htp_ops_context * octx;
|
||||
uint32_t npatches_per_thread; // patches = N*OH*OW (pure-DDR kernel)
|
||||
|
||||
uint32_t pe_rows_per_thread; // N*OH rows per worker
|
||||
uint32_t pe_src_row_bytes; // one output row's source: IC*KH*IW*4, rounded 256
|
||||
uint32_t pe_dst_row_bytes; // one output row's dst: OW*patch_stride*2, rounded 256
|
||||
|
||||
// Patch-embed DMA path VTCM ping-pong.
|
||||
uint8_t * pe_vtcm_src; // base of the 2x src buffers region
|
||||
uint8_t * pe_vtcm_dst; // base of the 2x dst buffers region
|
||||
uint32_t pe_src_size_per_thread; // 2 * pe_src_row_bytes
|
||||
uint32_t pe_dst_size_per_thread; // 2 * pe_dst_row_bytes
|
||||
};
|
||||
|
||||
// Per-op VTCM layout for the patch-embed DMA path
|
||||
struct htp_im2col_vtcm_layout {
|
||||
size_t off_src;
|
||||
size_t off_dst;
|
||||
size_t src_bytes_per_thread;
|
||||
size_t dst_bytes_per_thread;
|
||||
size_t total_bytes;
|
||||
};
|
||||
|
||||
static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * L,
|
||||
size_t src_row_bytes,
|
||||
size_t dst_row_bytes,
|
||||
uint32_t n_threads) {
|
||||
L->src_bytes_per_thread = 2 * src_row_bytes;
|
||||
L->dst_bytes_per_thread = 2 * dst_row_bytes;
|
||||
|
||||
L->off_src = 0;
|
||||
L->off_dst = L->off_src + L->src_bytes_per_thread * n_threads;
|
||||
L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
|
||||
}
|
||||
|
||||
#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \
|
||||
static void FNAME(unsigned int nth, unsigned int ith, void * data) { \
|
||||
struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \
|
||||
struct htp_ops_context * octx = ictx->octx; \
|
||||
struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \
|
||||
const struct htp_tensor * restrict src1 = octx->src[1]; \
|
||||
const struct htp_tensor * restrict dst = octx->dst; \
|
||||
const int32_t s0 = octx->op_params[0]; \
|
||||
const int32_t s1 = octx->op_params[1]; \
|
||||
const int32_t p0 = octx->op_params[2]; \
|
||||
const int32_t p1 = octx->op_params[3]; \
|
||||
const int32_t d0 = octx->op_params[4]; \
|
||||
const int32_t d1 = octx->op_params[5]; \
|
||||
const uint32_t N = src1->ne[3]; \
|
||||
const uint32_t IC = src1->ne[2]; \
|
||||
const uint32_t IH = src1->ne[1]; \
|
||||
const uint32_t IW = src1->ne[0]; \
|
||||
const uint32_t KH = octx->src[0]->ne[1]; \
|
||||
const uint32_t KW = octx->src[0]->ne[0]; \
|
||||
const uint32_t OH = dst->ne[2]; \
|
||||
const uint32_t OW = dst->ne[1]; \
|
||||
const uint32_t patch_stride = IC * KH * KW; \
|
||||
const float * restrict src_data = (const float *) src1->data; \
|
||||
DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \
|
||||
const uint32_t npatches = N * OH * OW; \
|
||||
const uint32_t patch_start = ictx->npatches_per_thread * ith; \
|
||||
const uint32_t patch_end = MIN(patch_start + ictx->npatches_per_thread, npatches); \
|
||||
if (patch_start >= patch_end) { \
|
||||
return; \
|
||||
} \
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \
|
||||
for (uint32_t p = patch_start; p < patch_end; p++) { \
|
||||
const uint32_t iow = p % OW; \
|
||||
const uint32_t ioh = (p / OW) % OH; \
|
||||
const uint32_t in = p / (OW * OH); \
|
||||
DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride; \
|
||||
for (uint32_t iic = 0; iic < IC; iic++) { \
|
||||
const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW; \
|
||||
for (uint32_t ikh = 0; ikh < KH; ikh++) { \
|
||||
const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \
|
||||
DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW; \
|
||||
if (iih < 0 || iih >= (int32_t) IH) { \
|
||||
SPLAT_FN(out_run, 0.0f, KW); \
|
||||
continue; \
|
||||
} \
|
||||
const int32_t iiw0 = (int32_t) iow * s0 - p0; \
|
||||
const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0; \
|
||||
if (d0 == 1) { \
|
||||
/* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */ \
|
||||
const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \
|
||||
int32_t hi = (int32_t) IW - iiw0; \
|
||||
if (hi > (int32_t) KW) { \
|
||||
hi = (int32_t) KW; \
|
||||
} \
|
||||
if (hi <= lo) { \
|
||||
SPLAT_FN(out_run, 0.0f, KW); \
|
||||
} else { \
|
||||
if (lo > 0) { \
|
||||
SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \
|
||||
} \
|
||||
COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo), \
|
||||
(uint32_t) (hi - lo)); \
|
||||
if (hi < (int32_t) KW) { \
|
||||
SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \
|
||||
} \
|
||||
} \
|
||||
continue; \
|
||||
} \
|
||||
for (uint32_t ikw = 0; ikw < KW; ikw++) { \
|
||||
const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0; \
|
||||
out_run[ikw] = (iiw < 0 || iiw >= (int32_t) IW) ? \
|
||||
(DST_CTYPE) 0.0f : \
|
||||
(DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw]; \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \
|
||||
}
|
||||
|
||||
IM2COL_PATCHEMBED_BODY(im2col_patchembed_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "f32-f16")
|
||||
IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "f32-f32")
|
||||
|
||||
#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \
|
||||
static void FNAME(unsigned int nth, unsigned int ith, void * data) { \
|
||||
struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \
|
||||
struct htp_ops_context * octx = ictx->octx; \
|
||||
struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \
|
||||
const struct htp_tensor * restrict src1 = octx->src[1]; \
|
||||
const struct htp_tensor * restrict dst = octx->dst; \
|
||||
const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0]; \
|
||||
const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; \
|
||||
const uint32_t OH = dst->ne[2], OW = dst->ne[1]; \
|
||||
const uint32_t patch_stride = IC * KH * KW; \
|
||||
const float * restrict src_data = (const float *) src1->data; \
|
||||
DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \
|
||||
dma_queue * dmaq = octx->ctx->dma[ith]; \
|
||||
uint8_t * src_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \
|
||||
uint8_t * dst_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \
|
||||
float * srcb = (float *) src_base; \
|
||||
DST_CTYPE * dstb = (DST_CTYPE *) dst_base; \
|
||||
const uint32_t nrows = N * OH; \
|
||||
const uint32_t per_thread = ictx->pe_rows_per_thread; \
|
||||
const uint32_t row_start = per_thread * ith; \
|
||||
const uint32_t row_end = MIN(row_start + per_thread, nrows); \
|
||||
if (row_start >= row_end) \
|
||||
return; \
|
||||
for (uint32_t r = row_start; r < row_end; r++) { \
|
||||
const uint32_t in = r / OH; \
|
||||
const uint32_t ioh = r % OH; \
|
||||
for (uint32_t ikh = 0; ikh < KH; ikh++) { \
|
||||
int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \
|
||||
int ok = (iih >= 0 && iih < (int32_t) IH); \
|
||||
for (uint32_t iic = 0; iic < IC; iic++) { \
|
||||
float * vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW; \
|
||||
const float * _vsrc = \
|
||||
ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst; \
|
||||
dma_queue_push_ddr_to_vtcm( \
|
||||
dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \
|
||||
IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0); \
|
||||
} \
|
||||
} \
|
||||
for (uint32_t i = 0; i < IC * KH; i++) \
|
||||
dma_queue_pop(dmaq); \
|
||||
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \
|
||||
for (uint32_t iow = 0; iow < OW; iow++) { \
|
||||
DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride; \
|
||||
for (uint32_t ikh = 0; ikh < KH; ikh++) { \
|
||||
int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \
|
||||
for (uint32_t iic = 0; iic < IC; iic++) { \
|
||||
DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \
|
||||
if (iih < 0 || iih >= (int32_t) IH) { \
|
||||
SPLAT_FN(out_run, 0.0f, KW); \
|
||||
continue; \
|
||||
} \
|
||||
const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW; \
|
||||
COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW); \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \
|
||||
DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride; \
|
||||
dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb), \
|
||||
OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1); \
|
||||
dma_queue_flush(dmaq); \
|
||||
} \
|
||||
}
|
||||
|
||||
IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16")
|
||||
IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "pe-dma-f32")
|
||||
|
||||
static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) {
|
||||
const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1];
|
||||
const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3];
|
||||
const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5];
|
||||
const int is_2D = octx->op_params[6] == 1;
|
||||
if (!is_2D) {
|
||||
return false;
|
||||
}
|
||||
if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
|
||||
return false;
|
||||
}
|
||||
const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
|
||||
if (s0 != (int32_t) KW || s1 != (int32_t) KH) {
|
||||
return false; // non-overlapping
|
||||
}
|
||||
if (p0 != 0 || p1 != 0) {
|
||||
return false; // no padding
|
||||
}
|
||||
if (d0 != 1 || d1 != 1) {
|
||||
return false; // no dilation
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// Sizes the per-thread 2x(src,dst) VTCM ping-pong for the patch-embed DMA path.
|
||||
// Returns false if it doesn't fit the VTCM budget (caller falls back).
|
||||
static bool im2col_patchembed_dma_fits(struct htp_ops_context * octx,
|
||||
struct htp_im2col_context * ictx,
|
||||
uint32_t n_threads) {
|
||||
const uint32_t IC = octx->src[1]->ne[2], IW = octx->src[1]->ne[0];
|
||||
const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
|
||||
const uint32_t OW = octx->dst->ne[1];
|
||||
const uint32_t patch_stride = IC * KH * KW;
|
||||
|
||||
ictx->pe_src_row_bytes = hex_round_up(IC * KH * IW * sizeof(float), 256);
|
||||
const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float);
|
||||
ictx->pe_dst_row_bytes = hex_round_up(OW * patch_stride * dst_elem, 256);
|
||||
|
||||
// 2 src + 2 dst buffers per thread (ping-pong), src region first then dst.
|
||||
struct htp_im2col_vtcm_layout L;
|
||||
htp_im2col_vtcm_layout_build(&L, ictx->pe_src_row_bytes, ictx->pe_dst_row_bytes, n_threads);
|
||||
if (L.total_bytes > octx->ctx->vtcm_size) {
|
||||
return false;
|
||||
}
|
||||
|
||||
uint8_t * const base = octx->ctx->vtcm_base;
|
||||
ictx->pe_vtcm_src = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src);
|
||||
ictx->pe_vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
|
||||
ictx->pe_src_size_per_thread = (uint32_t) L.src_bytes_per_thread;
|
||||
ictx->pe_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread;
|
||||
return true;
|
||||
}
|
||||
|
||||
int op_im2col(struct htp_ops_context * octx) {
|
||||
const struct htp_tensor * src1 = octx->src[1];
|
||||
const struct htp_tensor * dst = octx->dst;
|
||||
|
||||
if (src1->type != HTP_TYPE_F32 || (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32)) {
|
||||
FARF(ERROR, "im2col: only (F32 image -> F16/F32 columns) supported");
|
||||
return HTP_STATUS_NO_SUPPORT;
|
||||
}
|
||||
|
||||
const uint32_t N = src1->ne[3];
|
||||
const uint32_t OH = dst->ne[2];
|
||||
const uint32_t OW = dst->ne[1];
|
||||
const uint32_t npatches = N * OH * OW;
|
||||
const uint32_t n_threads = MIN(octx->n_threads, npatches);
|
||||
|
||||
if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) || n_threads == 0) {
|
||||
return HTP_STATUS_OK;
|
||||
}
|
||||
|
||||
struct htp_im2col_context ictx = { 0 };
|
||||
ictx.octx = octx;
|
||||
ictx.npatches_per_thread = (npatches + n_threads - 1) / n_threads;
|
||||
|
||||
// Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM);
|
||||
// everything else (padding/dilation/stride edges) -> pure-DDR kernel.
|
||||
if (im2col_use_patchembed_dma(octx)) {
|
||||
const uint32_t nrows = N * OH;
|
||||
const uint32_t pth = MIN(octx->n_threads, nrows);
|
||||
if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) {
|
||||
ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
|
||||
if (dst->type == HTP_TYPE_F16) {
|
||||
work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth);
|
||||
} else {
|
||||
work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_f32_thread, &ictx, pth);
|
||||
}
|
||||
return HTP_STATUS_OK;
|
||||
}
|
||||
// else: doesn't fit -> fall through to the pure-DDR kernel below.
|
||||
}
|
||||
|
||||
if (dst->type == HTP_TYPE_F16) {
|
||||
work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads);
|
||||
} else {
|
||||
work_queue_run(octx->ctx->work_queue, im2col_patchembed_f32_thread, &ictx, n_threads);
|
||||
}
|
||||
return HTP_STATUS_OK;
|
||||
}
|
||||
@@ -781,6 +781,9 @@ static int execute_op(struct htp_ops_context * octx) {
|
||||
case HTP_OP_PAD:
|
||||
return op_pad(octx);
|
||||
|
||||
case HTP_OP_IM2COL:
|
||||
return op_im2col(octx);
|
||||
|
||||
case HTP_OP_CONCAT:
|
||||
return op_concat(octx);
|
||||
|
||||
|
||||
@@ -114,10 +114,6 @@ if (GGML_HIP_NO_VMM)
|
||||
add_compile_definitions(GGML_HIP_NO_VMM)
|
||||
endif()
|
||||
|
||||
if (GGML_HIP_ROCWMMA_FATTN)
|
||||
add_compile_definitions(GGML_HIP_ROCWMMA_FATTN)
|
||||
endif()
|
||||
|
||||
if (NOT GGML_HIP_MMQ_MFMA)
|
||||
add_compile_definitions(GGML_HIP_NO_MMQ_MFMA)
|
||||
endif()
|
||||
@@ -158,5 +154,3 @@ if (GGML_HIP_RCCL)
|
||||
endif()
|
||||
|
||||
target_link_libraries(ggml-hip PRIVATE ggml-base hip::host roc::rocblas roc::hipblas)
|
||||
|
||||
target_compile_options(ggml-hip PRIVATE "$<$<COMPILE_LANGUAGE:HIP>:-ffast-math;-fno-finite-math-only>")
|
||||
|
||||
@@ -1252,6 +1252,21 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge(gg
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht(ggml_metal_library_t lib, int n) {
|
||||
char base[256];
|
||||
char name[256];
|
||||
|
||||
snprintf(base, 256, "kernel_fwht_f32_%d", n);
|
||||
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);
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
// note: reuse the argsort kernel for top_k
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal_library_t lib, const ggml_tensor * op) {
|
||||
assert(op->op == GGML_OP_TOP_K);
|
||||
|
||||
@@ -139,6 +139,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht (ggml_metal_library_t lib, int n);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse );
|
||||
|
||||
@@ -1157,6 +1157,10 @@ typedef struct {
|
||||
int32_t len;
|
||||
} ggml_metal_kargs_argsort_merge;
|
||||
|
||||
typedef struct {
|
||||
int32_t nrows;
|
||||
} ggml_metal_kargs_fwht;
|
||||
|
||||
typedef struct {
|
||||
int64_t ne0;
|
||||
float start;
|
||||
|
||||
@@ -1979,6 +1979,46 @@ int ggml_metal_op_pool_1d(ggml_metal_op_t ctx, int idx) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
// supported FWHT sizes, must stay in sync with the
|
||||
// kernel_fwht_f32_<N> templates in ggml-metal.metal
|
||||
static bool ggml_metal_fwht_supported_size(int64_t n) {
|
||||
return n == 64 || n == 128 || n == 256 || n == 512;
|
||||
}
|
||||
|
||||
int ggml_metal_op_fwht(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
ggml_metal_library_t lib = ctx->lib;
|
||||
ggml_metal_encoder_t enc = ctx->enc;
|
||||
|
||||
ggml_tensor * src1 = op->src[1];
|
||||
|
||||
const int64_t n = src1->ne[0];
|
||||
const int64_t nrows = ggml_nrows(src1);
|
||||
|
||||
ggml_metal_kargs_fwht args = {
|
||||
/*.nrows = */ (int32_t) nrows,
|
||||
};
|
||||
|
||||
auto pipeline = ggml_metal_library_get_pipeline_fwht(lib, n);
|
||||
|
||||
ggml_metal_encoder_set_pipeline(enc, pipeline);
|
||||
ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(src1), 1);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 2);
|
||||
|
||||
const int th_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline);
|
||||
const int simd_size = 32;
|
||||
|
||||
int sg_per_tg = 2;
|
||||
sg_per_tg = std::min(sg_per_tg, th_max/simd_size);
|
||||
sg_per_tg = std::max(sg_per_tg, 1);
|
||||
|
||||
const int64_t n_tg = (nrows + sg_per_tg - 1) / sg_per_tg;
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, n_tg, 1, 1, 32*sg_per_tg, 1, 1);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_pool_2d(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
@@ -2046,6 +2086,18 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_metal_library_t lib = ctx->lib;
|
||||
ggml_metal_encoder_t enc = ctx->enc;
|
||||
|
||||
const int32_t hint = ggml_get_op_params_i32(op, 1);
|
||||
|
||||
if (hint == GGML_HINT_SRC0_IS_HADAMARD) {
|
||||
if (op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->type == GGML_TYPE_F32 &&
|
||||
ggml_is_contiguous(op->src[1]) &&
|
||||
ggml_is_contiguous(op) &&
|
||||
ggml_are_same_shape(op->src[1], op) &&
|
||||
ggml_metal_fwht_supported_size(op->src[1]->ne[0])) {
|
||||
return ggml_metal_op_fwht(ctx, idx);
|
||||
}
|
||||
}
|
||||
const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev);
|
||||
|
||||
GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
|
||||
|
||||
@@ -64,6 +64,7 @@ int ggml_metal_op_set (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_cpy (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_pool_1d (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_pool_2d (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_fwht (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_mul_mat (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_mul_mat_id (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_add_id (ggml_metal_op_t ctx, int idx);
|
||||
|
||||
@@ -5762,7 +5762,7 @@ kernel void kernel_upscale_bicubic_f32(
|
||||
const float w_y2 = bicubic_weight1(1.0f - fd1);
|
||||
const float w_y3 = bicubic_weight2(2.0f - fd1);
|
||||
|
||||
const device const char * src_slice = src0 + i03 * args.nb03 + i02 * args.nb02;
|
||||
const device char * src_slice = src0 + i03 * args.nb03 + i02 * args.nb02;
|
||||
|
||||
device float * dst_ptr = (device float *)(dst + i3 * args.nb3 + i2 * args.nb2 + i1 * args.nb1);
|
||||
|
||||
@@ -6172,6 +6172,68 @@ kernel void kernel_argsort_merge_f32_i32(
|
||||
template [[host_name("kernel_argsort_merge_f32_i32_asc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_ASC>;
|
||||
template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_DESC>;
|
||||
|
||||
template<int N>
|
||||
kernel void kernel_fwht_f32(
|
||||
constant ggml_metal_kargs_fwht & args,
|
||||
device const float * src,
|
||||
device float * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
constexpr int NW = N_SIMDWIDTH;
|
||||
constexpr int NE = N / NW;
|
||||
|
||||
const float scale = 1.0f / sqrt((float) N);
|
||||
|
||||
const int sg_per_tg = ntg.x / NW;
|
||||
const int64_t r = tgpig.x * sg_per_tg + sgitg;
|
||||
if (r >= args.nrows) {
|
||||
return;
|
||||
}
|
||||
|
||||
src += r * N;
|
||||
dst += r * N;
|
||||
|
||||
const int lane = tiisg;
|
||||
|
||||
float reg[NE];
|
||||
for (int i = 0; i < NE; i++) {
|
||||
reg[i] = src[i*NW + lane]*scale;
|
||||
}
|
||||
for (int i = 1; i < NW; i *= 2) {
|
||||
for (int j = 0; j < NE; j++) {
|
||||
const float val = reg[j];
|
||||
const float val2 = simd_shuffle_xor(val, i);
|
||||
reg[j] = (lane & i) == 0 ? val2 + val : val2 - val;
|
||||
}
|
||||
}
|
||||
|
||||
for (int i = NW; i < N; i *= 2) {
|
||||
const int step = i / NW;
|
||||
for (int j = 0; j < NE; j += (2 * step)) {
|
||||
for (int k = 0; k < step; k++) {
|
||||
const float x = reg[j + k ];
|
||||
const float y = reg[j + k + step];
|
||||
reg[j + k] = x + y;
|
||||
reg[j + k + step] = x - y;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (int i = 0; i < NE; i++) {
|
||||
dst[i*NW + lane] = reg[i];
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_fwht_f32<64>) kernel_fwht_t;
|
||||
|
||||
template [[host_name("kernel_fwht_f32_64")]] kernel kernel_fwht_t kernel_fwht_f32<64>;
|
||||
template [[host_name("kernel_fwht_f32_128")]] kernel kernel_fwht_t kernel_fwht_f32<128>;
|
||||
template [[host_name("kernel_fwht_f32_256")]] kernel kernel_fwht_t kernel_fwht_f32<256>;
|
||||
template [[host_name("kernel_fwht_f32_512")]] kernel kernel_fwht_t kernel_fwht_f32<512>;
|
||||
|
||||
constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]];
|
||||
|
||||
constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]];
|
||||
|
||||
@@ -5,6 +5,8 @@ set(TARGET_NAME ggml-opencl)
|
||||
|
||||
ggml_add_backend_library(${TARGET_NAME}
|
||||
ggml-opencl.cpp
|
||||
cl-program-cache.cpp
|
||||
cl-program-cache.h
|
||||
../../include/ggml-opencl.h)
|
||||
target_link_libraries(${TARGET_NAME} PRIVATE ${OpenCL_LIBRARIES})
|
||||
target_include_directories(${TARGET_NAME} PRIVATE ${OpenCL_INCLUDE_DIRS})
|
||||
|
||||
@@ -0,0 +1,453 @@
|
||||
// Match the version setup ggml-opencl.cpp uses, so any cl.h declarations we
|
||||
// touch are consistent across this backend's translation units.
|
||||
#define CL_TARGET_OPENCL_VERSION GGML_OPENCL_TARGET_VERSION
|
||||
#define CL_USE_DEPRECATED_OPENCL_1_2_APIS
|
||||
|
||||
#include "cl-program-cache.h"
|
||||
|
||||
#include "ggml-impl.h" // GGML_LOG_INFO / WARN
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <system_error>
|
||||
#include <vector>
|
||||
|
||||
#if defined(_WIN32)
|
||||
# ifndef WIN32_LEAN_AND_MEAN
|
||||
# define WIN32_LEAN_AND_MEAN
|
||||
# endif
|
||||
# ifndef NOMINMAX
|
||||
# define NOMINMAX
|
||||
# endif
|
||||
# include <windows.h>
|
||||
# include <process.h>
|
||||
# define ggml_getpid() ((int) GetCurrentProcessId())
|
||||
#else
|
||||
# include <unistd.h>
|
||||
# define ggml_getpid() ((int) getpid())
|
||||
#endif
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// SHA-256 (FIPS 180-4). Self-contained, ~80 lines, public-domain reference.
|
||||
// Hot path is a few KB of source per kernel ⇒ <1 ms total per process init.
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
namespace {
|
||||
|
||||
struct sha256_ctx {
|
||||
uint32_t state[8];
|
||||
uint64_t bitlen;
|
||||
uint8_t buf[64];
|
||||
size_t buf_len;
|
||||
};
|
||||
|
||||
const uint32_t K256[64] = {
|
||||
0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5,
|
||||
0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174,
|
||||
0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da,
|
||||
0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967,
|
||||
0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85,
|
||||
0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070,
|
||||
0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3,
|
||||
0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2,
|
||||
};
|
||||
|
||||
inline uint32_t rotr32(uint32_t x, unsigned n) { return (x >> n) | (x << (32 - n)); }
|
||||
|
||||
void sha256_compress(uint32_t state[8], const uint8_t block[64]) {
|
||||
uint32_t w[64];
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
w[i] = ((uint32_t)block[i*4 ] << 24) |
|
||||
((uint32_t)block[i*4 + 1] << 16) |
|
||||
((uint32_t)block[i*4 + 2] << 8) |
|
||||
((uint32_t)block[i*4 + 3] );
|
||||
}
|
||||
for (int i = 16; i < 64; ++i) {
|
||||
uint32_t s0 = rotr32(w[i-15], 7) ^ rotr32(w[i-15], 18) ^ (w[i-15] >> 3);
|
||||
uint32_t s1 = rotr32(w[i-2], 17) ^ rotr32(w[i-2], 19) ^ (w[i-2] >> 10);
|
||||
w[i] = w[i-16] + s0 + w[i-7] + s1;
|
||||
}
|
||||
|
||||
uint32_t a = state[0],b = state[1],c = state[2],d = state[3],e = state[4],f = state[5],g = state[6],h = state[7];
|
||||
|
||||
for (int i = 0; i < 64; ++i) {
|
||||
uint32_t S1 = rotr32(e, 6) ^ rotr32(e, 11) ^ rotr32(e, 25);
|
||||
uint32_t ch = (e & f) ^ ((~e) & g);
|
||||
uint32_t t1 = h + S1 + ch + K256[i] + w[i];
|
||||
uint32_t S0 = rotr32(a, 2) ^ rotr32(a, 13) ^ rotr32(a, 22);
|
||||
uint32_t maj = (a & b) ^ (a & c) ^ (b & c);
|
||||
uint32_t t2 = S0 + maj;
|
||||
h = g; g = f; f = e; e = d + t1;
|
||||
d = c; c = b; b = a; a = t1 + t2;
|
||||
}
|
||||
state[0]+=a; state[1]+=b; state[2]+=c; state[3]+=d;
|
||||
state[4]+=e; state[5]+=f; state[6]+=g; state[7]+=h;
|
||||
}
|
||||
|
||||
void sha256_init(sha256_ctx & c) {
|
||||
c.state[0]=0x6a09e667; c.state[1]=0xbb67ae85; c.state[2]=0x3c6ef372; c.state[3]=0xa54ff53a;
|
||||
c.state[4]=0x510e527f; c.state[5]=0x9b05688c; c.state[6]=0x1f83d9ab; c.state[7]=0x5be0cd19;
|
||||
c.bitlen = 0;
|
||||
c.buf_len = 0;
|
||||
}
|
||||
|
||||
void sha256_update(sha256_ctx & c, const void * data, size_t len) {
|
||||
const uint8_t * p = (const uint8_t *) data;
|
||||
c.bitlen += (uint64_t) len * 8;
|
||||
if (c.buf_len > 0) {
|
||||
size_t n = 64 - c.buf_len;
|
||||
if (n > len) { n = len; }
|
||||
memcpy(c.buf + c.buf_len, p, n);
|
||||
c.buf_len += n;
|
||||
p += n;
|
||||
len -= n;
|
||||
if (c.buf_len == 64) {
|
||||
sha256_compress(c.state, c.buf);
|
||||
c.buf_len = 0;
|
||||
}
|
||||
}
|
||||
while (len >= 64) {
|
||||
sha256_compress(c.state, p);
|
||||
p += 64;
|
||||
len -= 64;
|
||||
}
|
||||
if (len > 0) {
|
||||
memcpy(c.buf, p, len);
|
||||
c.buf_len = len;
|
||||
}
|
||||
}
|
||||
|
||||
void sha256_final(sha256_ctx & c, uint8_t out[32]) {
|
||||
uint64_t bitlen = c.bitlen;
|
||||
c.buf[c.buf_len++] = 0x80;
|
||||
if (c.buf_len > 56) {
|
||||
while (c.buf_len < 64) { c.buf[c.buf_len++] = 0; }
|
||||
sha256_compress(c.state, c.buf);
|
||||
c.buf_len = 0;
|
||||
}
|
||||
while (c.buf_len < 56) { c.buf[c.buf_len++] = 0; }
|
||||
for (int i = 7; i >= 0; --i) { c.buf[c.buf_len++] = (uint8_t) (bitlen >> (i * 8)); }
|
||||
sha256_compress(c.state, c.buf);
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
out[i*4 ] = (uint8_t) (c.state[i] >> 24);
|
||||
out[i*4 + 1] = (uint8_t) (c.state[i] >> 16);
|
||||
out[i*4 + 2] = (uint8_t) (c.state[i] >> 8);
|
||||
out[i*4 + 3] = (uint8_t) (c.state[i] );
|
||||
}
|
||||
}
|
||||
|
||||
std::string sha256_hex(const uint8_t digest[32]) {
|
||||
static const char hex[] = "0123456789abcdef";
|
||||
std::string s(64, '0');
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
s[i*2 ] = hex[digest[i] >> 4];
|
||||
s[i*2 + 1] = hex[digest[i] & 0xf];
|
||||
}
|
||||
return s;
|
||||
}
|
||||
|
||||
std::string compute_key(const std::string & key_suffix,
|
||||
const char * source,
|
||||
const std::string & compile_opts) {
|
||||
sha256_ctx c;
|
||||
sha256_init(c);
|
||||
|
||||
static const uint8_t sep = 0;
|
||||
sha256_update(c, source, strlen(source));
|
||||
sha256_update(c, &sep, 1);
|
||||
sha256_update(c, compile_opts.data(), compile_opts.size());
|
||||
sha256_update(c, &sep, 1);
|
||||
sha256_update(c, key_suffix.data(), key_suffix.size());
|
||||
|
||||
uint8_t digest[32];
|
||||
sha256_final(c, digest);
|
||||
return sha256_hex(digest);
|
||||
}
|
||||
|
||||
bool make_dir_recursive(const std::string & path) {
|
||||
if (path.empty()) { return false; }
|
||||
// create_directories() already creates missing parents. It returns false
|
||||
// (with ec clear) when the directory is already there, so re-check.
|
||||
const fs::path p = fs::u8path(path);
|
||||
std::error_code ec;
|
||||
if (fs::create_directories(p, ec)) { return true; }
|
||||
std::error_code ec_stat;
|
||||
return fs::is_directory(p, ec_stat);
|
||||
}
|
||||
|
||||
std::string default_cache_dir() {
|
||||
#if defined(_WIN32)
|
||||
const char * base = std::getenv("LOCALAPPDATA");
|
||||
if (!base || !*base) { base = std::getenv("APPDATA"); }
|
||||
if (!base || !*base) { base = std::getenv("TEMP"); }
|
||||
if (!base || !*base) { base = "."; }
|
||||
return std::string(base) + "\\llama.cpp\\cl-cache";
|
||||
#elif defined(__APPLE__)
|
||||
const char * home = std::getenv("HOME");
|
||||
if (!home || !*home) { home = "."; }
|
||||
return std::string(home) + "/Library/Caches/llama.cpp/cl-cache";
|
||||
#else
|
||||
// The throwing overload aborts the process when no usable temp directory
|
||||
// exists (e.g. Android app contexts with TMPDIR unset); an empty return
|
||||
// here just disables the cache instead.
|
||||
std::error_code ec;
|
||||
const fs::path tmp_path = fs::temp_directory_path(ec);
|
||||
if (ec || tmp_path.empty()) { return {}; }
|
||||
return tmp_path.string() + "/llama.cpp/cl-cache";
|
||||
#endif
|
||||
}
|
||||
|
||||
// Query a NUL-terminated string from clGetDeviceInfo / clGetPlatformInfo.
|
||||
template <typename GetInfoFn, typename Object>
|
||||
std::string query_string(GetInfoFn fn, Object obj, cl_uint name) {
|
||||
size_t sz = 0;
|
||||
if (fn(obj, name, 0, nullptr, &sz) != CL_SUCCESS || sz == 0) {
|
||||
return {};
|
||||
}
|
||||
std::string s(sz, '\0');
|
||||
if (fn(obj, name, sz, &s[0], nullptr) != CL_SUCCESS) {
|
||||
return {};
|
||||
}
|
||||
if (!s.empty() && s.back() == '\0') {
|
||||
s.pop_back();
|
||||
}
|
||||
return s;
|
||||
}
|
||||
|
||||
std::string compute_key_suffix(cl_device_id device) {
|
||||
cl_platform_id platform = nullptr;
|
||||
clGetDeviceInfo(device, CL_DEVICE_PLATFORM, sizeof(platform), &platform, nullptr);
|
||||
|
||||
std::string s;
|
||||
s.reserve(512);
|
||||
s += query_string(clGetDeviceInfo, device, CL_DEVICE_NAME); s.push_back('\0');
|
||||
s += query_string(clGetDeviceInfo, device, CL_DRIVER_VERSION); s.push_back('\0');
|
||||
s += query_string(clGetDeviceInfo, device, CL_DEVICE_VERSION); s.push_back('\0');
|
||||
if (platform) {
|
||||
s += query_string(clGetPlatformInfo, platform, CL_PLATFORM_VERSION); s.push_back('\0');
|
||||
}
|
||||
s += "fmt=" + std::to_string(CL_PROGRAM_CACHE_FORMAT_VERSION);
|
||||
return s;
|
||||
}
|
||||
|
||||
const uint8_t MAGIC[8] = { 'G','G','M','L','C','L','B','C' };
|
||||
|
||||
bool read_all(const std::string & path, std::vector<uint8_t> & out) {
|
||||
std::ifstream f(fs::u8path(path), std::ios::binary);
|
||||
if (!f) { return false; }
|
||||
f.seekg(0, std::ios::end);
|
||||
std::streamsize sz = f.tellg();
|
||||
if (sz < 0) { return false; }
|
||||
f.seekg(0, std::ios::beg);
|
||||
out.resize((size_t) sz);
|
||||
if (sz > 0) { f.read((char *) out.data(), sz); }
|
||||
return f.good() || f.eof();
|
||||
}
|
||||
|
||||
bool write_atomic(const std::string & path, const uint8_t * data, size_t len) {
|
||||
const fs::path dst = fs::u8path(path);
|
||||
const fs::path tmp = fs::u8path(path + ".tmp." + std::to_string(ggml_getpid()));
|
||||
{
|
||||
std::ofstream f(tmp, std::ios::binary | std::ios::trunc);
|
||||
if (!f) { return false; }
|
||||
f.write((const char *) data, (std::streamsize) len);
|
||||
if (!f.good()) {
|
||||
std::error_code ec_rm;
|
||||
fs::remove(tmp, ec_rm);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
std::error_code ec;
|
||||
fs::rename(tmp, dst, ec);
|
||||
if (ec) {
|
||||
std::error_code ec_rm;
|
||||
fs::remove(tmp, ec_rm);
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
static bool cache_debug_enabled() {
|
||||
static int cached = -1;
|
||||
if (cached < 0) {
|
||||
const char * e = std::getenv("GGML_OPENCL_KERNEL_CACHE_DEBUG");
|
||||
cached = (e && *e) ? 1 : 0;
|
||||
}
|
||||
return cached != 0;
|
||||
}
|
||||
|
||||
static std::string opts_preview(const std::string & opts, size_t n = 120) {
|
||||
if (opts.size() <= n) { return opts; }
|
||||
return opts.substr(0, n) + "...";
|
||||
}
|
||||
|
||||
// Running cache tally (diagnostic; plain ints — a benign race in the rare
|
||||
// multi-threaded lazy-compile case at worst miscounts by one).
|
||||
static int g_cache_hits = 0, g_cache_misses = 0, g_cache_saves = 0;
|
||||
|
||||
// Debug trace directly to stderr
|
||||
static void cache_debug_line(const char * kind, const std::string & key,
|
||||
const char * source, const std::string & opts) {
|
||||
if (!cache_debug_enabled()) { return; }
|
||||
fprintf(stderr, "ggml_opencl: cache %-4s [h=%d m=%d s=%d] key=%s src=%zuB opts='%s'\n",
|
||||
kind, g_cache_hits, g_cache_misses, g_cache_saves,
|
||||
key.substr(0, 16).c_str(), strlen(source), opts_preview(opts).c_str());
|
||||
fflush(stderr);
|
||||
}
|
||||
|
||||
cl_program_cache_state cl_program_cache_init(cl_device_id device) {
|
||||
cl_program_cache_state st;
|
||||
|
||||
const char * env = std::getenv("GGML_OPENCL_KERNEL_CACHE_DIR");
|
||||
if (env && (!std::strcmp(env, "0") || !std::strcmp(env, "off") ||
|
||||
!std::strcmp(env, "none") || !std::strcmp(env, "disable") ||
|
||||
!std::strcmp(env, "disabled"))) {
|
||||
if (cache_debug_enabled()) {
|
||||
fprintf(stderr, "ggml_opencl: kernel cache disabled by GGML_OPENCL_KERNEL_CACHE_DIR=%s\n", env);
|
||||
fflush(stderr);
|
||||
}
|
||||
return st;
|
||||
}
|
||||
|
||||
std::string dir;
|
||||
if (!env || !*env || !std::strcmp(env, "1") || !std::strcmp(env, "default")) {
|
||||
dir = default_cache_dir();
|
||||
if (dir.empty()) {
|
||||
GGML_LOG_INFO("ggml_opencl: kernel cache disabled (no usable default cache directory)\n");
|
||||
return st;
|
||||
}
|
||||
} else {
|
||||
dir = env;
|
||||
}
|
||||
|
||||
if (!make_dir_recursive(dir)) {
|
||||
GGML_LOG_INFO("ggml_opencl: kernel cache disabled (cannot create directory '%s')\n", dir.c_str());
|
||||
return st;
|
||||
}
|
||||
|
||||
st.dir = dir;
|
||||
st.key_suffix = compute_key_suffix(device);
|
||||
GGML_LOG_INFO("ggml_opencl: kernel cache enabled at '%s'\n", st.dir.c_str());
|
||||
if (cache_debug_enabled()) {
|
||||
fprintf(stderr, "ggml_opencl: kernel cache enabled at '%s' "
|
||||
"(GGML_OPENCL_KERNEL_CACHE_DIR=off to disable)\n", st.dir.c_str());
|
||||
fflush(stderr);
|
||||
}
|
||||
return st;
|
||||
}
|
||||
|
||||
cl_program cl_program_cache_try_load(
|
||||
const cl_program_cache_state & state,
|
||||
cl_context context,
|
||||
cl_device_id device,
|
||||
const char * source,
|
||||
const std::string & compile_opts) {
|
||||
|
||||
if (state.dir.empty() || !source) { return nullptr; }
|
||||
|
||||
const std::string key = compute_key(state.key_suffix, source, compile_opts);
|
||||
const std::string path = state.dir + "/" + key + ".clbin";
|
||||
|
||||
std::vector<uint8_t> file;
|
||||
if (!read_all(path, file)) {
|
||||
++g_cache_misses;
|
||||
cache_debug_line("MISS", key, source, compile_opts);
|
||||
return nullptr;
|
||||
}
|
||||
if (file.size() < 16 || std::memcmp(file.data(), MAGIC, 8) != 0) { return nullptr; }
|
||||
|
||||
uint32_t fmt =
|
||||
((uint32_t) file[ 8]) | ((uint32_t) file[ 9] << 8) |
|
||||
((uint32_t) file[10] << 16) | ((uint32_t) file[11] << 24);
|
||||
if (fmt != CL_PROGRAM_CACHE_FORMAT_VERSION) { return nullptr; }
|
||||
|
||||
const size_t hdr_len = 16;
|
||||
const unsigned char * bin = file.data() + hdr_len;
|
||||
const size_t bin_len = file.size() - hdr_len;
|
||||
|
||||
cl_int err = CL_SUCCESS;
|
||||
cl_int bin_err = CL_SUCCESS;
|
||||
cl_program p = clCreateProgramWithBinary(context, 1, &device, &bin_len, &bin, &bin_err, &err);
|
||||
if (err != CL_SUCCESS || bin_err != CL_SUCCESS || p == nullptr) {
|
||||
if (p) { clReleaseProgram(p); }
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
err = clBuildProgram(p, 0, nullptr, compile_opts.c_str(), nullptr, nullptr);
|
||||
if (err != CL_SUCCESS) {
|
||||
clReleaseProgram(p);
|
||||
return nullptr;
|
||||
}
|
||||
++g_cache_hits;
|
||||
cache_debug_line("HIT", key, source, compile_opts);
|
||||
return p;
|
||||
}
|
||||
|
||||
void cl_program_cache_try_save(
|
||||
const cl_program_cache_state & state,
|
||||
cl_program program,
|
||||
cl_device_id /*device*/,
|
||||
const char * source,
|
||||
const std::string & compile_opts) {
|
||||
|
||||
if (state.dir.empty() || !program || !source) {
|
||||
return;
|
||||
}
|
||||
|
||||
cl_uint n_dev = 0;
|
||||
if (clGetProgramInfo(program, CL_PROGRAM_NUM_DEVICES, sizeof(n_dev), &n_dev, nullptr) != CL_SUCCESS || n_dev == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
std::vector<size_t> sizes(n_dev);
|
||||
if (clGetProgramInfo(program, CL_PROGRAM_BINARY_SIZES, sizeof(size_t) * n_dev, sizes.data(), nullptr) != CL_SUCCESS) {
|
||||
return;
|
||||
}
|
||||
if (sizes.empty() || sizes[0] == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
std::vector<std::vector<uint8_t>> binaries(n_dev);
|
||||
std::vector<unsigned char *> bin_ptrs(n_dev);
|
||||
for (cl_uint i = 0; i < n_dev; ++i) {
|
||||
binaries[i].resize(sizes[i]);
|
||||
bin_ptrs[i] = binaries[i].data();
|
||||
}
|
||||
if (clGetProgramInfo(program, CL_PROGRAM_BINARIES, sizeof(unsigned char *) * n_dev, bin_ptrs.data(), nullptr) != CL_SUCCESS) {
|
||||
return;
|
||||
}
|
||||
|
||||
// We only care about the first device's binary — that's the one we'd
|
||||
// re-load with on a future cache hit. Multi-device contexts aren't a
|
||||
// pattern this backend uses today.
|
||||
const std::vector<uint8_t> & bin = binaries[0];
|
||||
|
||||
std::vector<uint8_t> file;
|
||||
file.reserve(16 + bin.size());
|
||||
file.insert(file.end(), MAGIC, MAGIC + 8);
|
||||
uint32_t fmt = CL_PROGRAM_CACHE_FORMAT_VERSION;
|
||||
file.push_back((uint8_t) (fmt & 0xff));
|
||||
file.push_back((uint8_t) ((fmt >> 8) & 0xff));
|
||||
file.push_back((uint8_t) ((fmt >> 16) & 0xff));
|
||||
file.push_back((uint8_t) ((fmt >> 24) & 0xff));
|
||||
file.push_back(0); file.push_back(0); file.push_back(0); file.push_back(0); // reserved
|
||||
file.insert(file.end(), bin.begin(), bin.end());
|
||||
|
||||
const std::string key = compute_key(state.key_suffix, source, compile_opts);
|
||||
const std::string path = state.dir + "/" + key + ".clbin";
|
||||
if (!write_atomic(path, file.data(), file.size())) {
|
||||
GGML_LOG_INFO("ggml_opencl: kernel cache: failed to write '%s'\n", path.c_str());
|
||||
} else {
|
||||
++g_cache_saves;
|
||||
cache_debug_line("SAVE", key, source, compile_opts);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
// On-disk cache for OpenCL cl_program binaries. Lets a fresh process skip the
|
||||
// expensive clBuildProgram-from-source step when a binary for the exact same
|
||||
// (source, compile options, device, driver, platform) was previously saved.
|
||||
//
|
||||
// Activation: default on via GGML_OPENCL_KERNEL_CACHE_DIR:
|
||||
// unset / empty / "1" / "default" : platform default cache dir
|
||||
// (%LOCALAPPDATA%\llama.cpp\cl-cache,
|
||||
// ~/Library/Caches/llama.cpp/cl-cache,
|
||||
// <temp dir>/llama.cpp/cl-cache elsewhere)
|
||||
// "0" / "off" / "none" / "disable(d)" : disabled (all functions no-op)
|
||||
// any other value : used verbatim as the cache path
|
||||
// If the chosen directory cannot be created/used, the cache silently disables
|
||||
// itself for the process and falls back to source compile.
|
||||
// GGML_OPENCL_KERNEL_CACHE_DEBUG=1 prints a HIT/MISS/SAVE trace (with a running
|
||||
// tally) straight to stderr — visible even in tools that filter INFO/WARN logs;
|
||||
// redirect stderr to record it.
|
||||
//
|
||||
// Cache key (SHA-256 hex):
|
||||
// sha256(source_bytes || '\x00' ||
|
||||
// compile_opts || '\x00' ||
|
||||
// CL_DEVICE_NAME || '\x00' ||
|
||||
// CL_DRIVER_VERSION || '\x00' ||
|
||||
// CL_PLATFORM_VERSION || '\x00' ||
|
||||
// CL_PROGRAM_CACHE_FORMAT_VERSION)
|
||||
//
|
||||
// The key fully captures everything that can affect the produced binary,
|
||||
// without needing the host source revision (a kernel source change shows up
|
||||
// in source_bytes; a compile-option change shows up in compile_opts).
|
||||
//
|
||||
// File layout per cache entry: <cache_dir>/<sha256-hex>.clbin
|
||||
// bytes [0..7] : magic "GGMLCLBC"
|
||||
// bytes [8..11] : uint32_t format version (CL_PROGRAM_CACHE_FORMAT_VERSION)
|
||||
// bytes [12..15] : uint32_t reserved (0)
|
||||
// bytes [16..] : raw cl_program binary as returned by
|
||||
// clGetProgramInfo(CL_PROGRAM_BINARIES)
|
||||
//
|
||||
// Concurrency: writes go to <name>.tmp.<pid> then atomic rename. On race,
|
||||
// last-writer-wins. No locks.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <CL/cl.h>
|
||||
#include <string>
|
||||
|
||||
// Bumped manually if host-side OpenCL API usage changes in a way that
|
||||
// affects compile semantics but does not show up in source_bytes /
|
||||
// compile_opts (e.g. switching from clCreateProgramWithSource to
|
||||
// clCompileProgram + clLinkProgram, or changing how multiple sources
|
||||
// are concatenated). Most commits — including kernel changes — do NOT
|
||||
// require bumping this; the source bytes already capture those.
|
||||
#define CL_PROGRAM_CACHE_FORMAT_VERSION 1u
|
||||
|
||||
struct cl_program_cache_state {
|
||||
// Empty string means cache is disabled.
|
||||
std::string dir;
|
||||
// Concatenated device/driver/platform identity + cache format version,
|
||||
// computed once at init and folded into every key.
|
||||
std::string key_suffix;
|
||||
};
|
||||
|
||||
cl_program_cache_state cl_program_cache_init(cl_device_id device);
|
||||
|
||||
cl_program cl_program_cache_try_load(
|
||||
const cl_program_cache_state & state,
|
||||
cl_context context,
|
||||
cl_device_id device,
|
||||
const char * source,
|
||||
const std::string & compile_opts);
|
||||
|
||||
void cl_program_cache_try_save(
|
||||
const cl_program_cache_state & state,
|
||||
cl_program program,
|
||||
cl_device_id device,
|
||||
const char * source,
|
||||
const std::string & compile_opts);
|
||||
File diff suppressed because it is too large
Load Diff
@@ -199,9 +199,20 @@ if (GGML_SYCL_DEVICE_ARCH)
|
||||
-fsycl-targets=spir64_gen
|
||||
"SHELL:-Xsycl-target-backend=spir64_gen \"-device ${GGML_SYCL_DEVICE_ARCH}\""
|
||||
)
|
||||
|
||||
# Pass through parallel job (process) count for parallelising the
|
||||
# `llvm-foreach -- ocloc` invocation for compiling AOT device images.
|
||||
include(ProcessorCount)
|
||||
ProcessorCount(_ggml_sycl_nproc)
|
||||
if (_ggml_sycl_nproc LESS 1)
|
||||
set(_ggml_sycl_nproc 1)
|
||||
endif()
|
||||
set(GGML_SYCL_MAX_PARALLEL_LINK_JOBS ${_ggml_sycl_nproc} CACHE STRING
|
||||
"Parallel ocloc jobs for spir64_gen AOT device-image lowering")
|
||||
target_link_options(
|
||||
ggml-sycl PRIVATE
|
||||
-fsycl-targets=spir64_gen
|
||||
"SHELL:-Xsycl-target-backend=spir64_gen \"-device ${GGML_SYCL_DEVICE_ARCH}\""
|
||||
-fsycl-max-parallel-link-jobs=${GGML_SYCL_MAX_PARALLEL_LINK_JOBS}
|
||||
)
|
||||
endif()
|
||||
|
||||
@@ -65,6 +65,7 @@ extern int g_ggml_sycl_prioritize_dmmv;
|
||||
extern int g_ggml_sycl_enable_flash_attention;
|
||||
extern int g_ggml_sycl_dev2dev_memcpy;
|
||||
extern int g_ggml_sycl_fa_onednn;
|
||||
extern int g_ggml_sycl_fa_onednn_max_kv;
|
||||
|
||||
|
||||
#if defined(__clang__) && __has_builtin(__builtin_expect)
|
||||
|
||||
@@ -38,6 +38,12 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) {
|
||||
if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) {
|
||||
return false;
|
||||
}
|
||||
// Optional KV-length ceiling (GGML_SYCL_FA_ONEDNN_MAX_KV, 0 = unlimited). Escape hatch:
|
||||
// very long sequences make the fused SDPA slow enough to risk the xe driver watchdog on
|
||||
// some stacks; past the cap we fall back to the native FA kernel instead.
|
||||
if (g_ggml_sycl_fa_onednn_max_kv > 0 && K->ne[1] > g_ggml_sycl_fa_onednn_max_kv) {
|
||||
return false;
|
||||
}
|
||||
// gate for the following cases
|
||||
// 1. if the oneDNN graph Add node has no input --> skip
|
||||
// 2. types other than f16 need different logical_tensor declaration
|
||||
@@ -208,9 +214,17 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
|
||||
cont_to_f16_sycl<sycl::half>((const char *) V->data, Vf.get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream);
|
||||
|
||||
// divide-by-(1/scale) reproduces ggml's score *= kq_scale on the proven probe graph.
|
||||
//
|
||||
// The scale must not be uploaded with an async memcpy from a stack local: on the in-order
|
||||
// queue that copy waits behind the K/V staging kernels, and once those take long enough
|
||||
// (n_kv >= ~26k on B70) the host frame is recycled before the copy runs, feeding the SDPA a
|
||||
// garbage scale (output collapses to a repeated token). Write the scalar from a kernel
|
||||
// instead -- the value is captured into the command, so no host memory has to outlive the
|
||||
// call, and the enqueue stays async.
|
||||
const sycl::half scale_h = (sycl::half) (1.0f / kq_scale);
|
||||
ggml_sycl_pool_alloc<sycl::half> scbuf(ctx.pool(), 1);
|
||||
stream->memcpy(scbuf.get(), &scale_h, sizeof(sycl::half));
|
||||
sycl::half * const scale_dev = scbuf.get();
|
||||
stream->single_task([=]() { *scale_dev = scale_h; });
|
||||
|
||||
ggml_sycl_pool_alloc<sycl::half> outf(ctx.pool(), (size_t) H * q * d); // f16 contiguous SDPA out [mb,H,q,d]
|
||||
|
||||
@@ -232,7 +246,7 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
|
||||
if (r == E.id_q) return Qf.get();
|
||||
if (r == E.id_k) return Kf.get();
|
||||
if (r == E.id_v) return Vf.get();
|
||||
if (r == E.id_scale) return scbuf.get();
|
||||
if (r == E.id_scale) return scale_dev;
|
||||
if (r == E.id_mask) return (void *) mask->data;
|
||||
return nullptr;
|
||||
};
|
||||
@@ -245,14 +259,12 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
|
||||
E.cp.execute(strm, ti, {to});
|
||||
|
||||
permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream);
|
||||
// Single device: no sync is required, and actually PP perf is ~6% > wait_and_throw() (tested on llama-3.1-8b & qwen3.6-27b, both Q8_0, with Arc B70).
|
||||
// Any future multi-GPU refactor MUST re-measure this single-device path and keep the best
|
||||
// single-device PP speed. Otherwise (multiple devices/streams can race the reuse):
|
||||
// Single device needs no sync: the dnnl stream wraps this same in-order queue, so the SDPA
|
||||
// serializes with the staging kernels before it and the permute/pool reuse after it. The
|
||||
// garbage output formerly blamed on the missing sync here was the scale use-after-return
|
||||
// fixed above. Keep the conservative wait for multi-GPU, where other devices' streams can
|
||||
// race the pool:
|
||||
if (ggml_sycl_info().device_count > 1) {
|
||||
// cont_to_f16 -> oneDNN execute -> permute is async on this stream, but the
|
||||
// pool_alloc*s above free their device buffers at host return. Without this wait the next
|
||||
// scheduler op re-acquires those bytes while the GPU is still computing the SDPA, turning
|
||||
// it into garbage and collapsing multi-turn output to a single repeated token ("GGGGG...").
|
||||
stream->wait_and_throw();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -85,6 +85,7 @@ int g_ggml_sycl_enable_optimize = 1;
|
||||
int g_ggml_sycl_enable_graph = 0;
|
||||
int g_ggml_sycl_enable_dnn = 1;
|
||||
int g_ggml_sycl_fa_onednn = 1;
|
||||
int g_ggml_sycl_fa_onednn_max_kv = 0;
|
||||
int g_ggml_sycl_enable_vmm = 1;
|
||||
int g_ggml_sycl_enable_fusion = 1;
|
||||
int g_ggml_sycl_prioritize_dmmv = 0;
|
||||
@@ -287,6 +288,7 @@ static void ggml_check_sycl() try {
|
||||
g_ggml_sycl_enable_graph = ggml_sycl_get_env("GGML_SYCL_ENABLE_GRAPH", 0);
|
||||
g_ggml_sycl_enable_dnn = ggml_sycl_get_env("GGML_SYCL_ENABLE_DNN", 1);
|
||||
g_ggml_sycl_fa_onednn = ggml_sycl_get_env("GGML_SYCL_FA_ONEDNN", 1);
|
||||
g_ggml_sycl_fa_onednn_max_kv = ggml_sycl_get_env("GGML_SYCL_FA_ONEDNN_MAX_KV", 0);
|
||||
g_ggml_sycl_enable_vmm = ggml_sycl_get_env("GGML_SYCL_ENABLE_VMM", 1);
|
||||
g_ggml_sycl_enable_fusion = ggml_sycl_get_env("GGML_SYCL_ENABLE_FUSION", 1);
|
||||
g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0);
|
||||
@@ -359,6 +361,7 @@ static void ggml_check_sycl() try {
|
||||
GGML_LOG_INFO(" GGML_SYCL_ENABLE_DNN: DNN disabled by compile flag\n");
|
||||
GGML_LOG_INFO(" GGML_SYCL_FA_ONEDNN: %d\n", g_ggml_sycl_fa_onednn);
|
||||
#endif
|
||||
GGML_LOG_INFO(" GGML_SYCL_FA_ONEDNN_MAX_KV: %d\n", g_ggml_sycl_fa_onednn_max_kv);
|
||||
#ifdef SYCL_FLASH_ATTN
|
||||
GGML_LOG_INFO(" GGML_SYCL_ENABLE_FLASH_ATTN: %d\n", g_ggml_sycl_enable_flash_attention);
|
||||
#else
|
||||
|
||||
@@ -3490,7 +3490,7 @@ struct vk_fa_tuning_params {
|
||||
};
|
||||
|
||||
static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type);
|
||||
static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type = GGML_TYPE_F16);
|
||||
static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type = GGML_TYPE_F16, ggml_type v_type = GGML_TYPE_F16);
|
||||
|
||||
static vk_fa_tuning_params get_fa_tuning_params_scalar(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc) {
|
||||
|
||||
@@ -3646,7 +3646,7 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_
|
||||
bool shape_ok = (f32acc && device->coopmat_support_16x16x16_f32acc) ||
|
||||
(!f32acc && device->coopmat_support_16x16x16_f16acc);
|
||||
const vk_fa_tuning_params params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
|
||||
bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type);
|
||||
bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type, v_type);
|
||||
|
||||
if (!shape_ok || !shmem_ok) {
|
||||
path = FA_SCALAR;
|
||||
@@ -3658,11 +3658,6 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_
|
||||
path = FA_SCALAR;
|
||||
}
|
||||
|
||||
// Q1_0 K/V is only implemented on coopmat2 (flash_attn_cm2); there is no scalar FA shader for it.
|
||||
if ((k_type == GGML_TYPE_Q1_0 || v_type == GGML_TYPE_Q1_0) && device->coopmat2) {
|
||||
path = FA_COOPMAT2;
|
||||
}
|
||||
|
||||
switch (path) {
|
||||
case FA_SCALAR:
|
||||
return get_fa_tuning_params_scalar(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
|
||||
@@ -3904,16 +3899,27 @@ static uint32_t get_subgroup_size(const std::string &pipeline_name, const vk_dev
|
||||
return 0; // If no matching configuration is found
|
||||
}
|
||||
|
||||
// Whether scalar flash attention will use the MMQ path for the given k_type.
|
||||
static bool ggml_vk_fa_scalar_uses_mmq(const vk_device& device, ggml_type k_type) {
|
||||
// Whether scalar flash attention will use the MMQ path for the given K/V types.
|
||||
static bool ggml_vk_fa_type_needs_shmem(ggml_type type) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_IQ4_NL:
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
static bool ggml_vk_fa_scalar_uses_mmq(const vk_device& device, ggml_type k_type, ggml_type v_type) {
|
||||
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
|
||||
return device->integer_dot_product && device->subgroup_clustered &&
|
||||
!ggml_vk_fa_type_needs_shmem(v_type) &&
|
||||
(k_type == GGML_TYPE_Q4_0 || k_type == GGML_TYPE_Q4_1 ||
|
||||
k_type == GGML_TYPE_Q5_0 || k_type == GGML_TYPE_Q5_1 ||
|
||||
k_type == GGML_TYPE_Q8_0);
|
||||
#else
|
||||
GGML_UNUSED(device);
|
||||
GGML_UNUSED(k_type);
|
||||
GGML_UNUSED(v_type);
|
||||
return false;
|
||||
#endif
|
||||
}
|
||||
@@ -4246,7 +4252,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
const bool fa_ds = fa.first.subgroup_size == 0;
|
||||
|
||||
const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16;
|
||||
const bool use_mmq = ggml_vk_fa_scalar_uses_mmq(device, fa.first.k_type);
|
||||
const bool use_mmq = ggml_vk_fa_scalar_uses_mmq(device, fa.first.k_type, fa.first.v_type);
|
||||
const void * spv_data = nullptr;
|
||||
size_t spv_size = 0;
|
||||
const char *name = nullptr;
|
||||
@@ -10380,7 +10386,6 @@ static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx
|
||||
|
||||
static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type) {
|
||||
GGML_UNUSED(f32acc);
|
||||
GGML_UNUSED(v_type);
|
||||
// Needs to be kept up to date on shader changes
|
||||
const uint32_t wg_size = params.workgroup_size;
|
||||
const uint32_t Br = params.block_rows;
|
||||
@@ -10389,13 +10394,15 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con
|
||||
// BF16 uses the fp32 shader (FLOAT_TYPE=float)
|
||||
const uint32_t float_type_size = (device->fp16 && k_type != GGML_TYPE_BF16) ? sizeof(ggml_fp16_t) : sizeof(float);
|
||||
|
||||
const bool mmq = ggml_vk_fa_scalar_uses_mmq(device, k_type);
|
||||
const bool mmq = ggml_vk_fa_scalar_uses_mmq(device, k_type, v_type);
|
||||
|
||||
// tmpsh is overestimated slightly
|
||||
const uint32_t tmpsh = wg_size * sizeof(float);
|
||||
const uint32_t tmpshv4 = wg_size * 4 * float_type_size;
|
||||
|
||||
const uint32_t masksh = Bc * (Br + 1) * float_type_size;
|
||||
// DATA_A_IQ4_NL is compiled into the FA shaders unconditionally, so its shared table is always allocated.
|
||||
const uint32_t iq_shmem = 16 * float_type_size;
|
||||
|
||||
uint32_t Qf, kvsh, kblocksh_size;
|
||||
if (mmq) {
|
||||
@@ -10420,7 +10427,7 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con
|
||||
kblocksh_size = 0;
|
||||
}
|
||||
|
||||
const uint32_t total_size = tmpsh + tmpshv4 + masksh + Qf + kvsh + kblocksh_size;
|
||||
const uint32_t total_size = tmpsh + tmpshv4 + masksh + iq_shmem + Qf + kvsh + kblocksh_size;
|
||||
const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize;
|
||||
|
||||
VK_LOG_DEBUG("ggml_vk_flash_attn_scalar_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", mmq=" << mmq << ", total_size=" << total_size << ", supported=" << supported);
|
||||
@@ -10428,7 +10435,8 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con
|
||||
return supported;
|
||||
}
|
||||
|
||||
static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type) {
|
||||
static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type) {
|
||||
GGML_UNUSED(v_type);
|
||||
// Needs to be kept up to date on shader changes
|
||||
const uint32_t Br = params.block_rows;
|
||||
const uint32_t Bc = params.block_cols;
|
||||
@@ -10444,6 +10452,8 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co
|
||||
const uint32_t f16vec4 = 8;
|
||||
|
||||
const uint32_t tmpsh = (Bc / MatBc) * sizeof(float);
|
||||
// DATA_A_IQ4_NL is compiled into the FA shaders unconditionally, so its shared table is always allocated.
|
||||
const uint32_t iq_shmem = 16 * sizeof(ggml_fp16_t);
|
||||
|
||||
const uint32_t qstride = hsk_pad / 4 + 2;
|
||||
const uint32_t Qf = Br * qstride * f16vec4;
|
||||
@@ -10465,7 +10475,7 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co
|
||||
|
||||
const uint32_t slope = Br * acctype;
|
||||
|
||||
const uint32_t total_size = tmpsh + Qf + Psh + sfsh + ksh + pvsh + slope;
|
||||
const uint32_t total_size = tmpsh + iq_shmem + Qf + Psh + sfsh + ksh + pvsh + slope;
|
||||
const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize;
|
||||
|
||||
VK_LOG_DEBUG("ggml_vk_flash_attn_coopmat_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", f32acc=" << f32acc << ", total_size=" << total_size << ", supported=" << supported);
|
||||
@@ -17617,7 +17627,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
|
||||
if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) {
|
||||
return false;
|
||||
}
|
||||
auto fa_kv_ok = [coopmat2](ggml_type t) {
|
||||
auto fa_kv_ok = [](ggml_type t) {
|
||||
switch (t) {
|
||||
case GGML_TYPE_F32:
|
||||
case GGML_TYPE_F16:
|
||||
@@ -17627,9 +17637,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
|
||||
case GGML_TYPE_Q5_0:
|
||||
case GGML_TYPE_Q4_1:
|
||||
case GGML_TYPE_Q4_0:
|
||||
case GGML_TYPE_IQ4_NL:
|
||||
return true;
|
||||
case GGML_TYPE_Q1_0:
|
||||
return coopmat2;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -80,7 +80,9 @@ shared vec4 occupancy_limiter[LIMIT_OCCUPANCY_SHMEM > 0 ? LIMIT_OCCUPANCY_SHMEM
|
||||
|
||||
void main() {
|
||||
#ifdef NEEDS_INIT_IQ_SHMEM
|
||||
init_iq_shmem(gl_WorkGroupSize);
|
||||
if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) {
|
||||
init_iq_shmem(gl_WorkGroupSize);
|
||||
}
|
||||
#endif
|
||||
|
||||
init_indices();
|
||||
|
||||
@@ -97,8 +97,8 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
|
||||
#define FA_TYPE_Q5_0 6u
|
||||
#define FA_TYPE_Q5_1 7u
|
||||
#define FA_TYPE_Q8_0 8u
|
||||
#define FA_TYPE_IQ4_NL 20u
|
||||
#define FA_TYPE_BF16 30u
|
||||
#define FA_TYPE_Q1_0 41u
|
||||
|
||||
#if defined(BFLOAT16)
|
||||
#define O_TYPE float
|
||||
@@ -120,8 +120,8 @@ uint fa_block_elems(uint ty) {
|
||||
case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
|
||||
case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
|
||||
case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
|
||||
case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL);
|
||||
case FA_TYPE_BF16: return 1u;
|
||||
case FA_TYPE_Q1_0: return uint(QUANT_K_Q1_0); // cm2-only, harmless elsewhere
|
||||
default: return 1u;
|
||||
}
|
||||
}
|
||||
@@ -140,6 +140,13 @@ uint fa_quant_r_mmq(uint ty) {
|
||||
}
|
||||
}
|
||||
|
||||
bool fa_type_needs_shmem(uint ty) {
|
||||
switch (ty) {
|
||||
case FA_TYPE_IQ4_NL: return true;
|
||||
default: return false;
|
||||
}
|
||||
}
|
||||
|
||||
// These can't be `const` globals because GLSL forbids function calls in global
|
||||
// const initializers, even when the spec constants would let the driver fold
|
||||
// them. Macros expand at the use site and fold after specialization.
|
||||
|
||||
@@ -64,7 +64,9 @@ shared ACC_TYPE slope[Br];
|
||||
|
||||
void main() {
|
||||
#ifdef NEEDS_INIT_IQ_SHMEM
|
||||
init_iq_shmem(gl_WorkGroupSize);
|
||||
if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) {
|
||||
init_iq_shmem(gl_WorkGroupSize);
|
||||
}
|
||||
#endif
|
||||
|
||||
init_indices();
|
||||
|
||||
@@ -46,7 +46,7 @@ float16_t faDecodeK(const decodeBufFA_K bl_in, const uint blockCoords[2], const
|
||||
case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
|
||||
default: return float16_t(0);
|
||||
}
|
||||
}
|
||||
@@ -59,7 +59,7 @@ float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const
|
||||
case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
|
||||
default: return float16_t(0);
|
||||
}
|
||||
}
|
||||
@@ -67,26 +67,26 @@ float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const
|
||||
// V=4 vector decode for K/V; dispatches to per-format _v decoders.
|
||||
f16vec4 faDecodeKVector(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
|
||||
switch (FaTypeK) {
|
||||
case 0u: return f16vec4(decodeBufF32(bl_in).block);
|
||||
case 2u: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
|
||||
case 3u: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
|
||||
case 6u: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
|
||||
case 7u: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
|
||||
case 8u: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
|
||||
case 41u: return dequantFuncQ1_0_v(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
|
||||
case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
|
||||
default: return f16vec4(0);
|
||||
}
|
||||
}
|
||||
|
||||
f16vec4 faDecodeVVector(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
|
||||
switch (FaTypeV) {
|
||||
case 0u: return f16vec4(decodeBufF32(bl_in).block);
|
||||
case 2u: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
|
||||
case 3u: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
|
||||
case 6u: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
|
||||
case 7u: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
|
||||
case 8u: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
|
||||
case 41u: return dequantFuncQ1_0_v(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
|
||||
case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
|
||||
case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
|
||||
default: return f16vec4(0);
|
||||
}
|
||||
}
|
||||
@@ -169,6 +169,12 @@ ACC_TYPE perElemOpNonGqaSplitKStoreCol0(const in uint32_t r, const in uint32_t c
|
||||
}
|
||||
|
||||
void main() {
|
||||
#ifdef NEEDS_INIT_IQ_SHMEM
|
||||
if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) {
|
||||
init_iq_shmem(gl_WorkGroupSize);
|
||||
}
|
||||
#endif
|
||||
|
||||
init_indices();
|
||||
|
||||
tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutQ = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
|
||||
@@ -302,7 +308,7 @@ void main() {
|
||||
coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, HSK_pad, Bc, gl_MatrixUseB> K_T;
|
||||
|
||||
uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13;
|
||||
// F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Q4/Q8 family: bs_k==32. Q1_0: bs_k==128.
|
||||
// F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Quantized types: bs_k==32.
|
||||
#if defined(BFLOAT16)
|
||||
coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose);
|
||||
#else
|
||||
|
||||
@@ -27,6 +27,8 @@ layout (binding = 1) readonly buffer K_PACKED_Q5_1 { block_q5_1_packed16 data[];
|
||||
layout (binding = 2) readonly buffer V_PACKED_Q5_1 { block_q5_1_packed16 data[]; } v_packed_q5_1;
|
||||
layout (binding = 1) readonly buffer K_PACKED_Q8_0 { block_q8_0_packed16 data[]; } k_packed_q8_0;
|
||||
layout (binding = 2) readonly buffer V_PACKED_Q8_0 { block_q8_0_packed16 data[]; } v_packed_q8_0;
|
||||
layout (binding = 1) readonly buffer K_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } k_packed_iq4_nl;
|
||||
layout (binding = 2) readonly buffer V_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } v_packed_iq4_nl;
|
||||
|
||||
layout (binding = 1) readonly buffer K_PACKED_BF16 { u16vec4 data[]; } k_packed_bf16;
|
||||
layout (binding = 2) readonly buffer V_PACKED_BF16 { u16vec4 data[]; } v_packed_bf16;
|
||||
@@ -102,6 +104,17 @@ layout (binding = 1) readonly buffer K_PACKED_Q5_1_P32 { block_q5_1_packed32 dat
|
||||
return FLOAT_TYPE(BUF.data[a_offset + ib].d) * FLOAT_TYPEV4(v0.x, v0.y, v1.x, v1.y); \
|
||||
}
|
||||
|
||||
#define FA_DEQUANT4_IQ4_NL(BUF) { \
|
||||
const uint shift = (iqs & 0x10) >> 2; \
|
||||
const uint qs_i = (iqs & 0xC) >> 1; \
|
||||
const uint qsw = uint(BUF.data[a_offset + ib].qs[qs_i]) \
|
||||
| (uint(BUF.data[a_offset + ib].qs[qs_i + 1u]) << 16); \
|
||||
const FLOAT_TYPE d = FLOAT_TYPE(BUF.data[a_offset + ib].d); \
|
||||
const u8vec4 q = unpack8((qsw >> shift) & 0x0F0F0F0Fu); \
|
||||
return d * FLOAT_TYPEV4(kvalues_iq4nl[q.x], kvalues_iq4nl[q.y], \
|
||||
kvalues_iq4nl[q.z], kvalues_iq4nl[q.w]); \
|
||||
}
|
||||
|
||||
#define FA_DEQUANT4_BF16(BUF) \
|
||||
return FLOAT_TYPEV4(bf16_to_fp32(uvec4(BUF.data[(a_offset + ib) / 4])));
|
||||
|
||||
@@ -114,6 +127,7 @@ FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) {
|
||||
case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0)
|
||||
case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1)
|
||||
case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0)
|
||||
case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl)
|
||||
case FA_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16)
|
||||
}
|
||||
} else {
|
||||
@@ -124,6 +138,7 @@ FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) {
|
||||
case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0)
|
||||
case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1)
|
||||
case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0)
|
||||
case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl)
|
||||
case FA_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -673,6 +673,8 @@ void process_shaders() {
|
||||
fa_base_dict["ACC_TYPE"] = fp16 && f16acc ? "float16_t" : "float";
|
||||
fa_base_dict["ACC_TYPEV2"] = fp16 && f16acc ? "f16vec2" : "vec2";
|
||||
fa_base_dict["ACC_TYPEV4"] = fp16 && f16acc ? "f16vec4" : "vec4";
|
||||
// Compile IQ4_NL support into all FA variants so its shared LUT is available when K or V uses it.
|
||||
fa_base_dict["DATA_A_IQ4_NL"] = "1";
|
||||
if (fp16 && f16acc) {
|
||||
fa_base_dict["ACC_TYPE_MAX"] = "float16_t(65504.0)";
|
||||
}
|
||||
|
||||
+3
-1
@@ -7854,7 +7854,9 @@ void ggml_set_input(struct ggml_tensor * tensor) {
|
||||
}
|
||||
|
||||
void ggml_set_output(struct ggml_tensor * tensor) {
|
||||
tensor->flags |= GGML_TENSOR_FLAG_OUTPUT;
|
||||
for (struct ggml_tensor * cur = tensor; cur != NULL; cur = cur->view_src) {
|
||||
cur->flags |= GGML_TENSOR_FLAG_OUTPUT;
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_set_param(struct ggml_tensor * tensor) {
|
||||
|
||||
+1
-1
@@ -1424,7 +1424,7 @@ void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const vo
|
||||
struct gguf_writer_base {
|
||||
size_t written_bytes {0u};
|
||||
|
||||
~gguf_writer_base(void) = default;
|
||||
virtual ~gguf_writer_base(void) = default;
|
||||
|
||||
// we bet on devirtualization
|
||||
virtual void write(int8_t val) = 0;
|
||||
|
||||
+146
-1
@@ -145,6 +145,8 @@ class Keys:
|
||||
TOKEN_SHIFT_COUNT = "{arch}.token_shift_count"
|
||||
INTERLEAVE_MOE_LAYER_STEP = "{arch}.interleave_moe_layer_step"
|
||||
FULL_ATTENTION_INTERVAL = "{arch}.full_attention_interval"
|
||||
NUM_LOOPS = "{arch}.num_loops"
|
||||
SKIP_LOOP_FINAL_NORM = "{arch}.skip_loop_final_norm"
|
||||
HASH_LAYER_COUNT = "{arch}.hash_layer_count"
|
||||
ACTIVATION_SPARSITY_SCALE = "{arch}.activation_sparsity_scale"
|
||||
ALTUP_ACTIVE_IDX = "{arch}.altup.active_idx"
|
||||
@@ -159,6 +161,7 @@ class Keys:
|
||||
TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size"
|
||||
BLOCK_SIZE = "{arch}.block_size"
|
||||
NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual"
|
||||
NORM_BEFORE_FC = "{arch}.norm_before_fc"
|
||||
|
||||
class Attention:
|
||||
HEAD_COUNT = "{arch}.attention.head_count"
|
||||
@@ -200,6 +203,9 @@ class Keys:
|
||||
HEAD_COUNT = "{arch}.attention.indexer.head_count"
|
||||
KEY_LENGTH = "{arch}.attention.indexer.key_length"
|
||||
TOP_K = "{arch}.attention.indexer.top_k"
|
||||
BLOCK_SIZE = "{arch}.attention.indexer.block_size" # MSA
|
||||
LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks" # MSA
|
||||
TYPES = "{arch}.attention.indexer.types"
|
||||
|
||||
class HyperConnection:
|
||||
COUNT = "{arch}.hyper_connection.count"
|
||||
@@ -367,10 +373,17 @@ class Keys:
|
||||
FEED_FORWARD_LENGTH = "clip.audio.feed_forward_length"
|
||||
PROJECTION_DIM = "clip.audio.projection_dim"
|
||||
BLOCK_COUNT = "clip.audio.block_count"
|
||||
SUBSAMPLING_FACTOR = "clip.audio.subsampling_factor"
|
||||
CHUNK_SIZE = "clip.audio.chunk_size"
|
||||
CONV_KERNEL_SIZE = "clip.audio.conv_kernel_size"
|
||||
MAX_POS_EMB = "clip.audio.max_pos_emb"
|
||||
FEATURE_LAYERS = "clip.audio.feature_layer" # Granite Speech Plus
|
||||
RVQ_NUM_QUANTIZERS = "clip.audio.rvq.num_quantizers"
|
||||
RVQ_CODEBOOK_SIZE = "clip.audio.rvq.codebook_size"
|
||||
WA_PATTERN_MODE = "clip.audio.wa_pattern_mode" # per-layer -1 (full) / 0 (windowed)
|
||||
WINDOW_SIZE = "clip.audio.window_size"
|
||||
LOCAL_BLOCK_COUNT = "clip.audio.local_block_count" # mimo-v2.5: input_local_transformer layer count
|
||||
LOCAL_GROUP_SIZE = "clip.audio.local_group_size" # mimo-v2.5: input_local_transformer grouping size
|
||||
|
||||
class Attention:
|
||||
HEAD_COUNT = "clip.audio.attention.head_count"
|
||||
@@ -527,6 +540,7 @@ class MODEL_ARCH(IntEnum):
|
||||
APERTUS = auto()
|
||||
COGVLM = auto()
|
||||
MINIMAXM2 = auto()
|
||||
MINIMAXM3 = auto()
|
||||
RND1 = auto()
|
||||
PANGU_EMBED = auto()
|
||||
MISTRAL3 = auto()
|
||||
@@ -541,6 +555,7 @@ class MODEL_ARCH(IntEnum):
|
||||
KIMI_LINEAR = auto()
|
||||
TALKIE = auto()
|
||||
MELLUM = auto()
|
||||
NANBEIGE = auto()
|
||||
|
||||
|
||||
class VISION_PROJECTOR_TYPE(IntEnum):
|
||||
@@ -773,6 +788,9 @@ class MODEL_TENSOR(IntEnum):
|
||||
INDEXER_PROJ = auto()
|
||||
INDEXER_ATTN_K = auto()
|
||||
INDEXER_ATTN_Q_B = auto()
|
||||
INDEXER_Q_PROJ = auto()
|
||||
INDEXER_K_PROJ = auto()
|
||||
INDEXER_Q_NORM = auto()
|
||||
INDEXER_COMPRESSOR_WKV = auto()
|
||||
INDEXER_COMPRESSOR_WGATE = auto()
|
||||
INDEXER_COMPRESSOR_APE = auto()
|
||||
@@ -850,6 +868,8 @@ class MODEL_TENSOR(IntEnum):
|
||||
V_MM_UP = auto() # cogvlm
|
||||
V_MM_DOWN = auto() # cogvlm
|
||||
V_MM_GATE = auto() # cogvlm
|
||||
V_MM_MERGER_FC1 = auto() # minimax-m3 (patch-merge MLP)
|
||||
V_MM_MERGER_FC2 = auto() # minimax-m3 (patch-merge MLP)
|
||||
V_TOK_BOI = auto() # cogvlm
|
||||
V_TOK_EOI = auto() # cogvlm
|
||||
V_TOK_IMG_BEGIN = auto() # hunyuanvl
|
||||
@@ -933,6 +953,9 @@ class MODEL_TENSOR(IntEnum):
|
||||
A_ENC_FFN_SCALE_1 = auto() # gemma3n
|
||||
A_ENC_FFN_GATE_1 = auto() # lfm2, gemma3n
|
||||
A_ENC_FFN_DOWN_1 = auto() # lfm2, gemma3n
|
||||
A_ENC_DOWNSAMPLE_CONV = auto() # mimo-audio-tokenizer: post-transformer downsample conv
|
||||
A_ENC_DOWNSAMPLE_NORM = auto() # mimo-audio-tokenizer: post-transformer downsample norm
|
||||
A_ENC_RVQ_CODEBOOK = auto() # mimo-audio-tokenizer: residual vector quantizer codebook, per quantizer index
|
||||
A_MMPROJ = auto()
|
||||
A_MMPROJ_FC = auto()
|
||||
A_MM_NORM_PRE = auto()
|
||||
@@ -941,6 +964,17 @@ class MODEL_TENSOR(IntEnum):
|
||||
A_MM_HARD_EMB_NORM = auto() # gemma3n
|
||||
A_MM_SOFT_EMB_NORM = auto() # gemma3n
|
||||
A_MM_INP_PROJ = auto() # gemma3n
|
||||
A_MM_CODE_EMBD = auto() # mimo: text-side RVQ code embedding table ("text codebook"), merged 3D [n_channels, vocab, dim]
|
||||
A_MM_LOCAL_ATTN_Q = auto() # mimo: input_local_transformer (LLM-side connector)
|
||||
A_MM_LOCAL_ATTN_K = auto()
|
||||
A_MM_LOCAL_ATTN_V = auto()
|
||||
A_MM_LOCAL_ATTN_OUT = auto()
|
||||
A_MM_LOCAL_FFN_GATE = auto()
|
||||
A_MM_LOCAL_FFN_UP = auto()
|
||||
A_MM_LOCAL_FFN_DOWN = auto()
|
||||
A_MM_LOCAL_LN1 = auto()
|
||||
A_MM_LOCAL_LN2 = auto()
|
||||
A_MM_LOCAL_NORM = auto() # final norm after all input_local_transformer layers
|
||||
A_PER_DIM_K_SCALE = auto() # gemma4
|
||||
A_PER_DIM_SCALE = auto() # gemma4
|
||||
# nextn/mtp
|
||||
@@ -955,6 +989,10 @@ class MODEL_TENSOR(IntEnum):
|
||||
# eagle3
|
||||
FC = auto() # feature fusion layer
|
||||
D2T = auto() # draft to target vocabulary mapping
|
||||
# dspark
|
||||
DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed
|
||||
DSPARK_MARKOV_W2 = auto() # markov head: bias projection
|
||||
DSPARK_CONF_PROJ = auto() # confidence head
|
||||
# lfm2 audio
|
||||
A_ENC_NORM_CONV = auto()
|
||||
A_ENC_LINEAR_POS = auto()
|
||||
@@ -965,6 +1003,10 @@ class MODEL_TENSOR(IntEnum):
|
||||
A_ENC_CONV_NORM = auto() # SSM conv
|
||||
A_ENC_CONV_PW1 = auto()
|
||||
A_ENC_CONV_PW2 = auto()
|
||||
A_ENC_CONV_NORM_MEAN = auto() # parakeet
|
||||
A_ENC_CONV_NORM_VAR = auto() # parakeet
|
||||
A_ENC_MEL_FILTERS = auto() # parakeet
|
||||
A_ENC_WINDOW = auto() # parakeet
|
||||
A_CTC_OUT = auto()
|
||||
A_CTC_OUT_MID = auto()
|
||||
A_ENC_ATTN_REL_POS_EMB = auto()
|
||||
@@ -1109,6 +1151,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
|
||||
MODEL_ARCH.GROVEMOE: "grovemoe",
|
||||
MODEL_ARCH.APERTUS: "apertus",
|
||||
MODEL_ARCH.MINIMAXM2: "minimax-m2",
|
||||
MODEL_ARCH.MINIMAXM3: "minimax-m3",
|
||||
MODEL_ARCH.COGVLM: "cogvlm",
|
||||
MODEL_ARCH.RND1: "rnd1",
|
||||
MODEL_ARCH.PANGU_EMBED: "pangu-embedded",
|
||||
@@ -1124,6 +1167,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
|
||||
MODEL_ARCH.KIMI_LINEAR: "kimi-linear",
|
||||
MODEL_ARCH.TALKIE: "talkie",
|
||||
MODEL_ARCH.MELLUM: "mellum",
|
||||
MODEL_ARCH.NANBEIGE: "nanbeige",
|
||||
}
|
||||
|
||||
VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = {
|
||||
@@ -1354,6 +1398,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.indexer.proj",
|
||||
MODEL_TENSOR.INDEXER_ATTN_K: "blk.{bid}.indexer.attn_k",
|
||||
MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.indexer.attn_q_b",
|
||||
MODEL_TENSOR.INDEXER_Q_PROJ: "blk.{bid}.indexer.q_proj",
|
||||
MODEL_TENSOR.INDEXER_K_PROJ: "blk.{bid}.indexer.k_proj",
|
||||
MODEL_TENSOR.INDEXER_Q_NORM: "blk.{bid}.indexer.q_norm",
|
||||
MODEL_TENSOR.INDEXER_COMPRESSOR_WKV: "blk.{bid}.indexer_compressor_kv",
|
||||
MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: "blk.{bid}.indexer_compressor_gate",
|
||||
MODEL_TENSOR.INDEXER_COMPRESSOR_APE: "blk.{bid}.indexer_compressor_ape",
|
||||
@@ -1430,6 +1477,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.V_MM_UP: "mm.up",
|
||||
MODEL_TENSOR.V_MM_DOWN: "mm.down",
|
||||
MODEL_TENSOR.V_MM_GATE: "mm.gate",
|
||||
MODEL_TENSOR.V_MM_MERGER_FC1: "mm.merger.fc1",
|
||||
MODEL_TENSOR.V_MM_MERGER_FC2: "mm.merger.fc2",
|
||||
MODEL_TENSOR.V_TOK_BOI: "v.boi",
|
||||
MODEL_TENSOR.V_TOK_EOI: "v.eoi",
|
||||
MODEL_TENSOR.V_MM_PRE_NORM: "mm.pre_norm",
|
||||
@@ -1513,6 +1562,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.A_ENC_FFN_UP_1: "a.blk.{bid}.ffn_up_1",
|
||||
MODEL_TENSOR.A_ENC_FFN_GATE_1: "a.blk.{bid}.ffn_gate_1",
|
||||
MODEL_TENSOR.A_ENC_FFN_DOWN_1: "a.blk.{bid}.ffn_down_1",
|
||||
MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV: "a.downsample.conv",
|
||||
MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM: "a.downsample.norm",
|
||||
MODEL_TENSOR.A_ENC_RVQ_CODEBOOK: "a.rvq.codebook",
|
||||
MODEL_TENSOR.A_MMPROJ: "mm.a.mlp.{bid}",
|
||||
MODEL_TENSOR.A_MMPROJ_FC: "mm.a.fc",
|
||||
MODEL_TENSOR.A_MM_NORM_PRE: "mm.a.norm_pre",
|
||||
@@ -1521,6 +1573,17 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.A_MM_SOFT_EMB_NORM: "mm.a.soft_emb_norm", # gemma3n
|
||||
MODEL_TENSOR.A_MM_EMBEDDING: "mm.a.embedding", # gemma3n
|
||||
MODEL_TENSOR.A_MM_HARD_EMB_NORM: "mm.a.hard_emb_norm", # gemma3n
|
||||
MODEL_TENSOR.A_MM_CODE_EMBD: "mm.a.code_embd",
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_Q: "mm.a.local_blk.{bid}.attn_q",
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_K: "mm.a.local_blk.{bid}.attn_k",
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_V: "mm.a.local_blk.{bid}.attn_v",
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT: "mm.a.local_blk.{bid}.attn_out",
|
||||
MODEL_TENSOR.A_MM_LOCAL_FFN_GATE: "mm.a.local_blk.{bid}.ffn_gate",
|
||||
MODEL_TENSOR.A_MM_LOCAL_FFN_UP: "mm.a.local_blk.{bid}.ffn_up",
|
||||
MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN: "mm.a.local_blk.{bid}.ffn_down",
|
||||
MODEL_TENSOR.A_MM_LOCAL_LN1: "mm.a.local_blk.{bid}.ln1",
|
||||
MODEL_TENSOR.A_MM_LOCAL_LN2: "mm.a.local_blk.{bid}.ln2",
|
||||
MODEL_TENSOR.A_MM_LOCAL_NORM: "mm.a.local_norm",
|
||||
MODEL_TENSOR.A_PER_DIM_K_SCALE: "a.blk.{bid}.per_dim_k_scale", # gemma4
|
||||
MODEL_TENSOR.A_PER_DIM_SCALE: "a.blk.{bid}.per_dim_scale", # gemma4
|
||||
# lfm2 audio
|
||||
@@ -1533,6 +1596,10 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM: "a.blk.{bid}.conv_norm",
|
||||
MODEL_TENSOR.A_ENC_CONV_PW1: "a.blk.{bid}.conv_pw1",
|
||||
MODEL_TENSOR.A_ENC_CONV_PW2: "a.blk.{bid}.conv_pw2",
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_MEAN: "a.blk.{bid}.conv_norm_mean",
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_VAR: "a.blk.{bid}.conv_norm_var",
|
||||
MODEL_TENSOR.A_ENC_MEL_FILTERS: "a.mel_filters",
|
||||
MODEL_TENSOR.A_ENC_WINDOW: "a.window",
|
||||
MODEL_TENSOR.A_CTC_OUT: "a.enc_ctc_out",
|
||||
MODEL_TENSOR.A_CTC_OUT_MID: "a.enc_ctc_out_mid",
|
||||
MODEL_TENSOR.A_ENC_ATTN_REL_POS_EMB: "a.blk.{bid}.attn_rel_pos_emb",
|
||||
@@ -1563,6 +1630,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head",
|
||||
MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm",
|
||||
MODEL_TENSOR.FC: "fc",
|
||||
MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1",
|
||||
MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2",
|
||||
MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj",
|
||||
MODEL_TENSOR.D2T: "d2t",
|
||||
}
|
||||
|
||||
@@ -1626,6 +1696,8 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.V_RESMPL_QUERY,
|
||||
MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK,
|
||||
MODEL_TENSOR.V_MM_PATCH_MERGER,
|
||||
MODEL_TENSOR.V_MM_MERGER_FC1,
|
||||
MODEL_TENSOR.V_MM_MERGER_FC2,
|
||||
MODEL_TENSOR.V_DS_NORM,
|
||||
MODEL_TENSOR.V_DS_FC1,
|
||||
MODEL_TENSOR.V_DS_FC2,
|
||||
@@ -1720,10 +1792,24 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.A_ENC_FFN_UP_1,
|
||||
MODEL_TENSOR.A_ENC_FFN_GATE_1,
|
||||
MODEL_TENSOR.A_ENC_FFN_DOWN_1,
|
||||
MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV,
|
||||
MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM,
|
||||
MODEL_TENSOR.A_ENC_RVQ_CODEBOOK,
|
||||
MODEL_TENSOR.A_MMPROJ,
|
||||
MODEL_TENSOR.A_MMPROJ_FC,
|
||||
MODEL_TENSOR.A_MM_NORM_PRE,
|
||||
MODEL_TENSOR.A_MM_NORM_MID,
|
||||
MODEL_TENSOR.A_MM_CODE_EMBD,
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_Q,
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_K,
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_V,
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT,
|
||||
MODEL_TENSOR.A_MM_LOCAL_FFN_GATE,
|
||||
MODEL_TENSOR.A_MM_LOCAL_FFN_UP,
|
||||
MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN,
|
||||
MODEL_TENSOR.A_MM_LOCAL_LN1,
|
||||
MODEL_TENSOR.A_MM_LOCAL_LN2,
|
||||
MODEL_TENSOR.A_MM_LOCAL_NORM,
|
||||
MODEL_TENSOR.A_ENC_NORM_CONV,
|
||||
MODEL_TENSOR.A_ENC_LINEAR_POS,
|
||||
MODEL_TENSOR.A_ENC_POS_BIAS_U,
|
||||
@@ -1733,6 +1819,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM,
|
||||
MODEL_TENSOR.A_ENC_CONV_PW1,
|
||||
MODEL_TENSOR.A_ENC_CONV_PW2,
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_MEAN,
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_VAR,
|
||||
MODEL_TENSOR.A_ENC_MEL_FILTERS,
|
||||
MODEL_TENSOR.A_ENC_WINDOW,
|
||||
MODEL_TENSOR.A_MM_INP_PROJ,
|
||||
MODEL_TENSOR.A_MM_SOFT_EMB_NORM,
|
||||
MODEL_TENSOR.A_MM_EMBEDDING,
|
||||
@@ -4162,6 +4252,34 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.FFN_UP_EXP,
|
||||
MODEL_TENSOR.FFN_EXP_PROBS_B,
|
||||
],
|
||||
MODEL_ARCH.MINIMAXM3: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
MODEL_TENSOR.OUTPUT,
|
||||
MODEL_TENSOR.ATTN_NORM,
|
||||
MODEL_TENSOR.ATTN_Q,
|
||||
MODEL_TENSOR.ATTN_Q_NORM,
|
||||
MODEL_TENSOR.ATTN_K,
|
||||
MODEL_TENSOR.ATTN_K_NORM,
|
||||
MODEL_TENSOR.ATTN_V,
|
||||
MODEL_TENSOR.ATTN_OUT,
|
||||
MODEL_TENSOR.FFN_NORM,
|
||||
MODEL_TENSOR.FFN_GATE_INP,
|
||||
MODEL_TENSOR.FFN_EXP_PROBS_B,
|
||||
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_TENSOR.FFN_GATE,
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
MODEL_TENSOR.INDEXER_Q_PROJ,
|
||||
MODEL_TENSOR.INDEXER_K_PROJ,
|
||||
MODEL_TENSOR.INDEXER_Q_NORM,
|
||||
MODEL_TENSOR.INDEXER_K_NORM,
|
||||
],
|
||||
MODEL_ARCH.COGVLM: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
@@ -4246,6 +4364,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
MODEL_TENSOR.FC,
|
||||
MODEL_TENSOR.ENC_OUTPUT_NORM,
|
||||
MODEL_TENSOR.D2T,
|
||||
],
|
||||
MODEL_ARCH.DFLASH: [
|
||||
@@ -4263,6 +4382,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
MODEL_TENSOR.FC,
|
||||
MODEL_TENSOR.ENC_OUTPUT_NORM,
|
||||
# optional DSpark heads
|
||||
MODEL_TENSOR.DSPARK_MARKOV_W1,
|
||||
MODEL_TENSOR.DSPARK_MARKOV_W2,
|
||||
MODEL_TENSOR.DSPARK_CONF_PROJ,
|
||||
],
|
||||
MODEL_ARCH.MISTRAL4: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
@@ -4460,7 +4583,22 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.FFN_DOWN_EXP,
|
||||
MODEL_TENSOR.FFN_UP_EXP,
|
||||
],
|
||||
# TODO
|
||||
MODEL_ARCH.NANBEIGE: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
MODEL_TENSOR.OUTPUT,
|
||||
MODEL_TENSOR.ROPE_FREQS,
|
||||
MODEL_TENSOR.ATTN_NORM,
|
||||
MODEL_TENSOR.ATTN_Q,
|
||||
MODEL_TENSOR.ATTN_K,
|
||||
MODEL_TENSOR.ATTN_V,
|
||||
MODEL_TENSOR.ATTN_OUT,
|
||||
MODEL_TENSOR.ATTN_ROT_EMBD,
|
||||
MODEL_TENSOR.FFN_NORM,
|
||||
MODEL_TENSOR.FFN_GATE,
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
],
|
||||
}
|
||||
|
||||
# tensors that will not be serialized
|
||||
@@ -4527,6 +4665,10 @@ MODEL_TENSOR_SKIP: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.ROPE_FREQS,
|
||||
MODEL_TENSOR.ATTN_ROT_EMBD,
|
||||
],
|
||||
MODEL_ARCH.NANBEIGE: [
|
||||
MODEL_TENSOR.ROPE_FREQS,
|
||||
MODEL_TENSOR.ATTN_ROT_EMBD,
|
||||
],
|
||||
}
|
||||
|
||||
#
|
||||
@@ -4732,9 +4874,12 @@ class VisionProjectorType:
|
||||
YOUTUVL = "youtuvl"
|
||||
NEMOTRON_V2_VL = "nemotron_v2_vl"
|
||||
HUNYUANVL = "hunyuanvl"
|
||||
PARAKEET = "parakeet" # audio
|
||||
MINIMAXM3 = "minimax_m3"
|
||||
MINICPMV4_6 = "minicpmv4_6"
|
||||
GRANITE_SPEECH = "granite_speech" # audio
|
||||
MIMOVL = "mimovl"
|
||||
MIMO_AUDIO = "mimo_audio"
|
||||
GRANITE4_VISION = "granite4_vision"
|
||||
|
||||
|
||||
|
||||
@@ -793,6 +793,16 @@ class GGUFWriter:
|
||||
def add_indexer_top_k(self, top_k: int) -> None:
|
||||
self.add_uint32(Keys.Attention.Indexer.TOP_K.format(arch=self.arch), top_k)
|
||||
|
||||
def add_indexer_block_size(self, block_size: int) -> None:
|
||||
self.add_uint32(Keys.Attention.Indexer.BLOCK_SIZE.format(arch=self.arch), block_size)
|
||||
|
||||
def add_indexer_local_blocks(self, local_blocks: int) -> None:
|
||||
self.add_uint32(Keys.Attention.Indexer.LOCAL_BLOCKS.format(arch=self.arch), local_blocks)
|
||||
|
||||
def add_indexer_types(self, value: Sequence[bool]) -> None:
|
||||
key = Keys.Attention.Indexer.TYPES.format(arch=self.arch)
|
||||
self.add_array(key, value)
|
||||
|
||||
def add_max_alibi_bias(self, bias: float) -> None:
|
||||
self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias)
|
||||
|
||||
@@ -898,6 +908,12 @@ class GGUFWriter:
|
||||
def add_token_shift_count(self, count: int) -> None:
|
||||
self.add_uint32(Keys.LLM.TOKEN_SHIFT_COUNT.format(arch=self.arch), count)
|
||||
|
||||
def add_num_loops(self, count: int) -> None:
|
||||
self.add_uint32(Keys.LLM.NUM_LOOPS.format(arch=self.arch), count)
|
||||
|
||||
def add_skip_loop_final_norm(self, value: bool) -> None:
|
||||
self.add_bool(Keys.LLM.SKIP_LOOP_FINAL_NORM.format(arch=self.arch), value)
|
||||
|
||||
def add_interleave_moe_layer_step(self, value: int) -> None:
|
||||
self.add_uint32(Keys.LLM.INTERLEAVE_MOE_LAYER_STEP.format(arch=self.arch), value)
|
||||
|
||||
@@ -955,6 +971,9 @@ class GGUFWriter:
|
||||
def add_norm_before_residual(self, value: bool) -> None:
|
||||
self.add_bool(Keys.LLM.NORM_BEFORE_RESIDUAL.format(arch=self.arch), value)
|
||||
|
||||
def add_norm_before_fc(self, value: bool) -> None:
|
||||
self.add_bool(Keys.LLM.NORM_BEFORE_FC.format(arch=self.arch), value)
|
||||
|
||||
def add_attention_output_group_count(self, count: int) -> None:
|
||||
self.add_uint32(Keys.Attention.OUTPUT_GROUP_COUNT.format(arch=self.arch), count)
|
||||
|
||||
@@ -1334,9 +1353,30 @@ class GGUFWriter:
|
||||
def add_audio_num_mel_bins(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.NUM_MEL_BINS, value)
|
||||
|
||||
def add_audio_rvq_num_quantizers(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.RVQ_NUM_QUANTIZERS, value)
|
||||
|
||||
def add_audio_rvq_codebook_size(self, values: Sequence[int]) -> None:
|
||||
self.add_array(Keys.ClipAudio.RVQ_CODEBOOK_SIZE, values)
|
||||
|
||||
def add_audio_wa_pattern_mode(self, modes: Sequence[int]) -> None:
|
||||
self.add_array(Keys.ClipAudio.WA_PATTERN_MODE, modes)
|
||||
|
||||
def add_audio_window_size(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.WINDOW_SIZE, value)
|
||||
|
||||
def add_audio_local_block_count(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.LOCAL_BLOCK_COUNT, value)
|
||||
|
||||
def add_audio_local_group_size(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.LOCAL_GROUP_SIZE, value)
|
||||
|
||||
def add_audio_stack_factor(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.Projector.STACK_FACTOR, value)
|
||||
|
||||
def add_audio_subsampling_factor(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.SUBSAMPLING_FACTOR, value)
|
||||
|
||||
def add_audio_chunk_size(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.CHUNK_SIZE, value)
|
||||
|
||||
|
||||
@@ -1264,7 +1264,8 @@ class TensorNameMap:
|
||||
),
|
||||
|
||||
MODEL_TENSOR.INDEXER_K_NORM: (
|
||||
"model.layers.{bid}.self_attn.indexer.k_norm", # DSA
|
||||
"model.layers.{bid}.self_attn.indexer.k_norm", # DSA
|
||||
"model.layers.{bid}.self_attn.index_k_norm", # MSA
|
||||
),
|
||||
|
||||
MODEL_TENSOR.INDEXER_PROJ: (
|
||||
@@ -1279,6 +1280,18 @@ class TensorNameMap:
|
||||
"model.layers.{bid}.self_attn.indexer.wq_b", # DSA
|
||||
),
|
||||
|
||||
MODEL_TENSOR.INDEXER_Q_PROJ: (
|
||||
"model.layers.{bid}.self_attn.index_q_proj", # MSA
|
||||
),
|
||||
|
||||
MODEL_TENSOR.INDEXER_K_PROJ: (
|
||||
"model.layers.{bid}.self_attn.index_k_proj", # MSA
|
||||
),
|
||||
|
||||
MODEL_TENSOR.INDEXER_Q_NORM: (
|
||||
"model.layers.{bid}.self_attn.index_q_norm", # MSA
|
||||
),
|
||||
|
||||
############################################################################
|
||||
# TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg
|
||||
MODEL_TENSOR.ENC_OUTPUT_NORM: (
|
||||
@@ -1291,6 +1304,18 @@ class TensorNameMap:
|
||||
"model.fc", # dflash
|
||||
),
|
||||
|
||||
MODEL_TENSOR.DSPARK_MARKOV_W1: (
|
||||
"model.markov_head.markov_w1", # dspark
|
||||
),
|
||||
|
||||
MODEL_TENSOR.DSPARK_MARKOV_W2: (
|
||||
"model.markov_head.markov_w2", # dspark
|
||||
),
|
||||
|
||||
MODEL_TENSOR.DSPARK_CONF_PROJ: (
|
||||
"model.confidence_head.proj", # dspark
|
||||
),
|
||||
|
||||
MODEL_TENSOR.CLS: (
|
||||
"classifier", # jina
|
||||
"classifier.dense", # roberta
|
||||
@@ -1825,6 +1850,14 @@ class TensorNameMap:
|
||||
"visual.downsample", # glm4v
|
||||
),
|
||||
|
||||
MODEL_TENSOR.V_MM_MERGER_FC1: (
|
||||
"patch_merge_mlp.linear_1", # minimax-m3
|
||||
),
|
||||
|
||||
MODEL_TENSOR.V_MM_MERGER_FC2: (
|
||||
"patch_merge_mlp.linear_2", # minimax-m3
|
||||
),
|
||||
|
||||
MODEL_TENSOR.V_DS_NORM: (
|
||||
"model.visual.deepstack_merger_list.{bid}.norm", # deepstack in qwen3vl
|
||||
),
|
||||
@@ -2074,6 +2107,8 @@ class TensorNameMap:
|
||||
"conformer.pre_encode.conv.{bid}", # lfm2
|
||||
"model.audio_tower.subsample_conv_projection.conv_{bid}.conv", # gemma3n
|
||||
"conformer.subsample_conv_projection.layer{bid}.conv", # gemma4
|
||||
"sound_encoder.encoder.subsampling.layers.{bid}", # parakeet
|
||||
"encoder.conv{bid}", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV1D_NORM: (
|
||||
@@ -2098,6 +2133,7 @@ class TensorNameMap:
|
||||
MODEL_TENSOR.A_POST_NORM: (
|
||||
"audio_tower.layer_norm", # ultravox
|
||||
"audio_tower.ln_post", # qwen2omni
|
||||
"encoder.layer_norm", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_ATTN_Q: (
|
||||
@@ -2105,7 +2141,9 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.self_attn.linear_q", # lfm2
|
||||
"conformer.layers.{bid}.attention.attn.q_proj", # gemma3n
|
||||
"conformer.layers.{bid}.self_attn.q_proj", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.q_proj", # parakeet
|
||||
"encoder.layers.{bid}.attn.to_q", # granite_speech
|
||||
"encoder.layers.{bid}.self_attn.q_proj", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_ATTN_K: (
|
||||
@@ -2113,7 +2151,9 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.self_attn.linear_k", # lfm2
|
||||
"conformer.layers.{bid}.attention.attn.k_proj", # gemma3n
|
||||
"conformer.layers.{bid}.self_attn.k_proj", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.k_proj", # parakeet
|
||||
"encoder.layers.{bid}.attn.to_k", # granite_speech (split from to_kv)
|
||||
"encoder.layers.{bid}.self_attn.k_proj", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_ATTN_V: (
|
||||
@@ -2121,7 +2161,9 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.self_attn.linear_v", # lfm2
|
||||
"conformer.layers.{bid}.attention.attn.v_proj", # gemma3n
|
||||
"conformer.layers.{bid}.self_attn.v_proj", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.v_proj", # parakeet
|
||||
"encoder.layers.{bid}.attn.to_v", # granite_speech (split from to_kv)
|
||||
"encoder.layers.{bid}.self_attn.v_proj", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_ATTN_K_REL: (
|
||||
@@ -2149,7 +2191,9 @@ class TensorNameMap:
|
||||
"audio_tower.layers.{bid}.self_attn_layer_norm", # ultravox
|
||||
"conformer.layers.{bid}.norm_self_att", # lfm2
|
||||
"conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.norm_self_att", # parakeet
|
||||
"encoder.layers.{bid}.attn.pre_norm", # granite_speech
|
||||
"encoder.layers.{bid}.self_attn_layer_norm", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_OUTPUT: (
|
||||
@@ -2157,20 +2201,25 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.self_attn.linear_out", # lfm2
|
||||
"conformer.layers.{bid}.attention.post", # gemma3n
|
||||
"conformer.layers.{bid}.self_attn.post", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.o_proj", # parakeet
|
||||
"encoder.layers.{bid}.attn.to_out", # granite_speech
|
||||
"encoder.layers.{bid}.self_attn.out_proj", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_OUTPUT_NORM: (
|
||||
"audio_tower.layers.{bid}.final_layer_norm", # ultravox
|
||||
"conformer.layers.{bid}.norm_out", # lfm2
|
||||
"conformer.layers.{bid}.attention.post_norm", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.norm_out", # parakeet
|
||||
"encoder.layers.{bid}.post_norm", # granite_speech
|
||||
"encoder.layers.{bid}.final_layer_norm", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_FFN_NORM: (
|
||||
"conformer.layers.{bid}.norm_feed_forward1", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_start.pre_layer_norm", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward1.pre_layer_norm", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.norm_feed_forward1", # parakeet
|
||||
"encoder.layers.{bid}.ff1.pre_norm", # granite_speech
|
||||
),
|
||||
|
||||
@@ -2188,7 +2237,9 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.feed_forward1.linear1", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.feed_forward1.linear1", # parakeet
|
||||
"encoder.layers.{bid}.ff1.up_proj", # granite_speech
|
||||
"encoder.layers.{bid}.fc1", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_FFN_GATE: (),
|
||||
@@ -2198,13 +2249,16 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.feed_forward1.linear2", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.feed_forward1.linear2", # parakeet
|
||||
"encoder.layers.{bid}.ff1.down_proj", # granite_speech
|
||||
"encoder.layers.{bid}.fc2", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_FFN_UP_1: (
|
||||
"conformer.layers.{bid}.feed_forward2.linear1", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_end.ffw_layer_1", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward2.ffw_layer_1", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.feed_forward2.linear1", # parakeet
|
||||
"encoder.layers.{bid}.ff2.up_proj", # granite_speech
|
||||
),
|
||||
|
||||
@@ -2212,6 +2266,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.feed_forward2.linear2", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_end.ffw_layer_2", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward2.ffw_layer_2", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.feed_forward2.linear2", # parakeet
|
||||
"encoder.layers.{bid}.ff2.down_proj", # granite_speech
|
||||
),
|
||||
|
||||
@@ -2219,9 +2274,23 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.norm_feed_forward2", # lfm2
|
||||
"conformer.layers.{bid}.ffw_layer_end.pre_layer_norm", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward2.pre_layer_norm", # gemma4
|
||||
"sound_encoder.encoder.layers.{bid}.norm_feed_forward2", # parakeet
|
||||
"encoder.layers.{bid}.ff2.pre_norm", # granite_speech
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV: (
|
||||
"encoder.down_sample_layer.0", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM: (
|
||||
"encoder.down_sample_norm", # mimo-audio-tokenizer
|
||||
),
|
||||
|
||||
# note: the raw per-quantizer "encoder.quantizer.vq.layers.{i}._codebook.embed"
|
||||
# tensors are merged (padded + stacked, like MoE experts) into this single 3D
|
||||
# tensor in conversion code, so no raw-name mapping is registered here.
|
||||
MODEL_TENSOR.A_ENC_RVQ_CODEBOOK: (),
|
||||
|
||||
MODEL_TENSOR.A_ENC_FFN_POST_NORM_1: (
|
||||
"conformer.layers.{bid}.ffw_layer_end.post_layer_norm", # gemma3n
|
||||
"conformer.layers.{bid}.feed_forward2.post_layer_norm", # gemma4
|
||||
@@ -2234,20 +2303,24 @@ class TensorNameMap:
|
||||
MODEL_TENSOR.A_ENC_LINEAR_POS: (
|
||||
"conformer.layers.{bid}.self_attn.linear_pos", # lfm2
|
||||
"conformer.layers.{bid}.attention.attn.relative_position_embedding.pos_proj", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.relative_k_proj", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_POS_BIAS_U: (
|
||||
"conformer.layers.{bid}.self_attn.pos_bias_u", # lfm2
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.bias_u", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_POS_BIAS_V: (
|
||||
"conformer.layers.{bid}.self_attn.pos_bias_v", # lfm2
|
||||
"sound_encoder.encoder.layers.{bid}.self_attn.bias_v", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_OUT: (
|
||||
"conformer.pre_encode.out", # lfm2
|
||||
"model.audio_tower.subsample_conv_projection.input_proj_linear", # gemma3n (note: it should be A_ENC_INP_PROJ, this is a mistake; it should be corrected in C++ code when it's supported)
|
||||
"conformer.output_proj", # gemma4
|
||||
"sound_encoder.encoder.subsampling.linear", # parakeet
|
||||
),
|
||||
|
||||
# note: some tensors below has "audio." pseudo-prefix, to prevent conflicts with vision tensors
|
||||
@@ -2257,6 +2330,7 @@ class TensorNameMap:
|
||||
"audio.multi_modal_projector.linear_{bid}", # ultravox, meralion
|
||||
"audio_adapter.model.{bid}", # lfm2
|
||||
"audio_tower.proj{bid}", # qwen3omni
|
||||
"sound_projection.linear{bid}", # parakeet (linear1, linear2)
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_MMPROJ_FC: (
|
||||
@@ -2267,39 +2341,89 @@ class TensorNameMap:
|
||||
|
||||
MODEL_TENSOR.A_MM_NORM_PRE: (
|
||||
"audio.multi_modal_projector.ln_pre", # ultravox
|
||||
"sound_projection.norm", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_MM_NORM_MID: (
|
||||
"audio.multi_modal_projector.ln_mid", # ultravox
|
||||
),
|
||||
|
||||
# note: the raw per-channel "speech_embeddings.{i}" tensors are merged
|
||||
# (stacked, like MoE experts) into this single 3D tensor in conversion
|
||||
# code, so no raw-name mapping is registered here.
|
||||
MODEL_TENSOR.A_MM_CODE_EMBD: (),
|
||||
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_Q: (
|
||||
"audio_encoder.input_local_transformer.layers.{bid}.self_attn.q_proj", # mimo-v2.5
|
||||
),
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_K: (
|
||||
"audio_encoder.input_local_transformer.layers.{bid}.self_attn.k_proj", # mimo-v2.5
|
||||
),
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_V: (
|
||||
"audio_encoder.input_local_transformer.layers.{bid}.self_attn.v_proj", # mimo-v2.5
|
||||
),
|
||||
MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT: (
|
||||
"audio_encoder.input_local_transformer.layers.{bid}.self_attn.o_proj", # mimo-v2.5
|
||||
),
|
||||
MODEL_TENSOR.A_MM_LOCAL_FFN_GATE: (
|
||||
"audio_encoder.input_local_transformer.layers.{bid}.mlp.gate_proj", # mimo-v2.5
|
||||
),
|
||||
MODEL_TENSOR.A_MM_LOCAL_FFN_UP: (
|
||||
"audio_encoder.input_local_transformer.layers.{bid}.mlp.up_proj", # mimo-v2.5
|
||||
),
|
||||
MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN: (
|
||||
"audio_encoder.input_local_transformer.layers.{bid}.mlp.down_proj", # mimo-v2.5
|
||||
),
|
||||
MODEL_TENSOR.A_MM_LOCAL_LN1: (
|
||||
"audio_encoder.input_local_transformer.layers.{bid}.input_layernorm", # mimo-v2.5
|
||||
),
|
||||
MODEL_TENSOR.A_MM_LOCAL_LN2: (
|
||||
"audio_encoder.input_local_transformer.layers.{bid}.post_attention_layernorm", # mimo-v2.5
|
||||
),
|
||||
MODEL_TENSOR.A_MM_LOCAL_NORM: (
|
||||
"audio_encoder.input_local_transformer.norm", # mimo-v2.5
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_DW: (
|
||||
"conformer.layers.{bid}.conv.depthwise_conv", # lfm2
|
||||
"conformer.layers.{bid}.lconv1d.depthwise_conv1d", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.conv.depthwise_conv", # parakeet
|
||||
"encoder.layers.{bid}.conv.depth_conv.conv", # granite_speech
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM: (
|
||||
"conformer.layers.{bid}.conv.batch_norm", # lfm2
|
||||
"conformer.layers.{bid}.lconv1d.pre_layer_norm", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.conv.norm", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_MEAN: (
|
||||
"sound_encoder.encoder.layers.{bid}.conv.norm.running_mean", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_VAR: (
|
||||
"sound_encoder.encoder.layers.{bid}.conv.norm.running_var", # parakeet
|
||||
"encoder.layers.{bid}.conv.batch_norm", # granite_speech
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_PW1: (
|
||||
"conformer.layers.{bid}.conv.pointwise_conv1", # lfm2
|
||||
"conformer.layers.{bid}.lconv1d.linear_start", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.conv.pointwise_conv1", # parakeet
|
||||
"encoder.layers.{bid}.conv.up_conv", # granite_speech
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_PW2: (
|
||||
"conformer.layers.{bid}.conv.pointwise_conv2", # lfm2
|
||||
"conformer.layers.{bid}.lconv1d.linear_end", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.conv.pointwise_conv2", # parakeet
|
||||
"encoder.layers.{bid}.conv.down_conv", # granite_speech
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_NORM_CONV: (
|
||||
"conformer.layers.{bid}.norm_conv", # lfm2
|
||||
"conformer.layers.{bid}.lconv1d.conv_norm", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.norm_conv", # parakeet
|
||||
"encoder.layers.{bid}.conv.norm", # granite_speech
|
||||
),
|
||||
|
||||
@@ -2311,6 +2435,14 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.attention.attn.per_dim_scale", # gemma4
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_MEL_FILTERS: (
|
||||
"sound_encoder.encoder.feature_extractor.featurizer.fb", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_WINDOW: (
|
||||
"sound_encoder.encoder.feature_extractor.featurizer.window", # parakeet
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_MM_EMBEDDING: (
|
||||
"model.embed_audio.embedding", # gemma3n
|
||||
),
|
||||
|
||||
+5
-4
@@ -203,10 +203,11 @@ extern "C" {
|
||||
};
|
||||
|
||||
enum llama_load_mode {
|
||||
LLAMA_LOAD_MODE_NONE = 0, // no special loading mode
|
||||
LLAMA_LOAD_MODE_MMAP = 1, // memory map the model
|
||||
LLAMA_LOAD_MODE_MLOCK = 2, // mmap + force system to keep model in RAM rather than swapping or compressing
|
||||
LLAMA_LOAD_MODE_DIRECT_IO = 3, // use direct I/O if available
|
||||
LLAMA_LOAD_MODE_NONE = 0, // no special loading mode
|
||||
LLAMA_LOAD_MODE_MMAP = 1, // memory map the model
|
||||
LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing
|
||||
LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing
|
||||
LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available
|
||||
};
|
||||
|
||||
LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode);
|
||||
|
||||
@@ -0,0 +1,247 @@
|
||||
{# ---------- special token variables ---------- #}
|
||||
{%- set ns_token = ']<]minimax[>[' -%}
|
||||
{%- set bod_token = ']~!b[' -%}
|
||||
{%- set bos_token = ']~b]' -%}
|
||||
{%- set eos_token = '[e~[' -%}
|
||||
{%- set toolcall_begin_token = ns_token ~ '<tool_call>' -%}
|
||||
{%- set toolcall_end_token = ns_token ~ '</tool_call>' -%}
|
||||
{%- set think_begin_token = '<mm:think>' -%}
|
||||
{%- set think_end_token = '</mm:think>' -%}
|
||||
{%- set image_token = ']<]image[>[' -%}
|
||||
{%- set video_token = ']<]video[>[' -%}
|
||||
{#- Thinking mode: "enabled" / "disabled" / "adaptive" / not defined -#}
|
||||
{#- Recursive XML renderer for tool_call arguments ======================== -#}
|
||||
{#- None values are intentionally skipped in mapping iteration so that
|
||||
`<key>null</key>` (which would round-trip to the literal string "null")
|
||||
never appears in the rendered tool_call. The convention is: omit the
|
||||
field entirely. The top-level `_args` loop applies the same rule.
|
||||
The `val is none` branch below is a safety net only — upstream cleaning
|
||||
(drop_none_in_tool_arguments) should ensure no None ever reaches here. -#}
|
||||
{%- macro to_xml(val, ns) -%}
|
||||
{%- if val is mapping -%}
|
||||
{%- for k, v in val.items() if v is not none -%}
|
||||
{{ ns }}<{{ k }}>{{ to_xml(v, ns) }}{{ ns }}</{{ k }}>
|
||||
{%- endfor -%}
|
||||
{%- elif val is iterable and val is not string -%}
|
||||
{%- for item in val -%}
|
||||
{{ ns }}<item>{{ to_xml(item, ns) }}{{ ns }}</item>
|
||||
{%- endfor -%}
|
||||
{%- elif val is none -%}
|
||||
{#- Should be unreachable when upstream cleaning is applied. -#}
|
||||
{%- elif val is boolean -%}
|
||||
{{ val | tojson }}
|
||||
{%- else -%}
|
||||
{{ val }}
|
||||
{%- endif -%}
|
||||
{%- endmacro -%}
|
||||
{#- Tool Rendering Functions ============================================== -#}
|
||||
{%- macro render_tool_namespace(namespace_name, tool_list) -%}
|
||||
{%- for tool in tool_list -%}
|
||||
<tool>{{ tool.function | tojson(ensure_ascii=False) }}</tool>
|
||||
{% endfor -%}
|
||||
{%- endmacro -%}
|
||||
{%- macro visible_text(content) -%}
|
||||
{%- if content is string -%}
|
||||
{{ content }}
|
||||
{%- elif content is iterable and content is not mapping -%}
|
||||
{%- for item in content -%}
|
||||
{%- if item is mapping and item.type == 'text' -%}
|
||||
{{- item.text }}
|
||||
{%- elif item is mapping and item.type == 'image' -%}
|
||||
{{- image_token }}
|
||||
{%- elif item is mapping and item.type == 'video' -%}
|
||||
{{- video_token}}
|
||||
{%- elif item is string -%}
|
||||
{{- item }}
|
||||
{%- endif -%}
|
||||
{%- endfor -%}
|
||||
{%- elif content is none -%}
|
||||
{{- '' }}
|
||||
{%- else -%}
|
||||
{{- content }}
|
||||
{%- endif -%}
|
||||
{%- endmacro -%}
|
||||
{#- System Message Construction ============================================ -#}
|
||||
{%- macro build_system_message(system_message) -%}
|
||||
{%- if system_message and system_message.content -%}
|
||||
{{- visible_text(system_message.content) }}
|
||||
{%- else -%}
|
||||
{{- 'Your model version is MiniMax-M3, developed by MiniMax. Knowledge cutoff: January 2026. Founded in early 2022, MiniMax is a global AI foundation model company committed to advancing the frontiers of AI towards AGI.' }}
|
||||
{%- endif -%}
|
||||
|
||||
{#- Thinking mode instructions -#}
|
||||
{{- '\n\n<thinking_instructions>\n' }}
|
||||
{{- 'You have a thinking capability that allows you to reason step by step before responding. When thinking is enabled, wrap your reasoning in ' ~ think_begin_token ~ think_end_token ~ ' tags before your response. When thinking is disabled, begin your response directly after the ' ~ think_end_token ~ ' prefix. When thinking is adaptive, decide on your own whether to think for the current turn.\n' }}
|
||||
{%- if thinking_mode is defined -%}
|
||||
{%- if thinking_mode == "enabled" -%}
|
||||
{{- 'Current thinking mode: enabled. You MUST think step by step before every response, including after receiving function/tool results.\n' }}
|
||||
{%- elif thinking_mode == "disabled" -%}
|
||||
{{- 'Current thinking mode: disabled. Do not output any thinking process.\n' }}
|
||||
{%- elif thinking_mode == "adaptive" -%}
|
||||
{{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }}
|
||||
{%- endif -%}
|
||||
{%- else -%}
|
||||
{{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }}
|
||||
{%- endif -%}
|
||||
{{- '</thinking_instructions>' }}
|
||||
{%- endmacro -%}
|
||||
{%- macro build_developer_message(developer_message) -%}
|
||||
{%- if developer_message and developer_message.content -%}
|
||||
{{- visible_text(developer_message.content) }}
|
||||
{%- else -%}
|
||||
{%- if model_identity is not defined -%}
|
||||
{%- set model_identity = "You are a helpful assistant." -%}
|
||||
{%- endif -%}
|
||||
{{- model_identity }}
|
||||
{%- endif -%}
|
||||
{%- endmacro -%}
|
||||
{#- Main Template Logic ================================================= -#}
|
||||
{#- Role mapping: root -> system sp (high priority), system/developer -> developer sp (low priority) -#}
|
||||
{%- set system_message = none -%}
|
||||
{%- set developer_message = none -%}
|
||||
{%- set conversation_messages = messages -%}
|
||||
{%- if messages and messages[0].role == "root" -%}
|
||||
{%- set system_message = messages[0] -%}
|
||||
{%- set conversation_messages = messages[1:] -%}
|
||||
{%- if conversation_messages and conversation_messages[0].role in ["system", "developer"] -%}
|
||||
{%- set developer_message = conversation_messages[0] -%}
|
||||
{%- set conversation_messages = conversation_messages[1:] -%}
|
||||
{%- endif -%}
|
||||
{%- elif messages and messages[0].role in ["system", "developer"] -%}
|
||||
{%- set developer_message = messages[0] -%}
|
||||
{%- set conversation_messages = messages[1:] -%}
|
||||
{%- endif -%}
|
||||
{#- Render system sp (higher priority, root role only) -#}
|
||||
{{- bod_token ~ bos_token ~ 'system' ~ '\n' }}
|
||||
{{- build_system_message(system_message) }}
|
||||
{{- eos_token ~ '\n' }}
|
||||
|
||||
{#- Render developer sp (lower priority: system/developer role + tools) -#}
|
||||
{{- bos_token ~ 'developer' ~ '\n' }}
|
||||
{{- build_developer_message(developer_message) }}
|
||||
{%- if tools -%}
|
||||
{{- '\n\n' ~ '# Tools' ~ '\n' ~ 'You may call one or more tools to assist with the user query.\nHere are the tools available in JSONSchema format:' ~ '\n' }}
|
||||
{{- '\n' ~ '<tools>' ~ '\n' }}
|
||||
{{- render_tool_namespace("functions", tools) }}
|
||||
{{- '</tools>' ~ '\n\n' }}
|
||||
{{- 'To call tools, wrap all invocations in a single ' ~ toolcall_begin_token ~ toolcall_end_token ~ ' block. Parameter values containing nested objects or arrays are recursively expanded into XML elements. Example:\n' }}
|
||||
{{- '\n' ~ toolcall_begin_token ~ '\n' }}
|
||||
{{- ns_token + '<invoke name="tool-name-1">' }}
|
||||
{{- ns_token + '<param-1>value-1' + ns_token + '</param-1>' }}
|
||||
{{- ns_token + '<param-2>' }}
|
||||
{{- ns_token + '<item>' }}
|
||||
{{- ns_token + '<key-a>val-a' + ns_token + '</key-a>' }}
|
||||
{{- ns_token + '<key-b>val-b' + ns_token + '</key-b>' }}
|
||||
{{- ns_token + '</item>' }}
|
||||
{{- ns_token + '</param-2>' }}
|
||||
{{- ns_token + '</invoke>\n' }}
|
||||
{{- ns_token + '<invoke name="tool-name-2">' }}
|
||||
{{- ns_token + '<param-1>value-1' + ns_token + '</param-1>' }}
|
||||
{{- ns_token + '</invoke>\n' }}
|
||||
{{- toolcall_end_token }}
|
||||
{%- endif -%}
|
||||
{{- eos_token ~ '\n' }}
|
||||
|
||||
{#- Render messages -#}
|
||||
{%- set last_tool_call = namespace(name=none) -%}
|
||||
{%- for message in conversation_messages -%}
|
||||
{%- if message.role == 'assistant' -%}
|
||||
{{- bos_token ~ 'ai' ~ '\n' }}
|
||||
|
||||
{%- set reasoning_content = '' %}
|
||||
{%- set content = visible_text(message.content) %}
|
||||
{%- if message.reasoning_content is string %}
|
||||
{%- set reasoning_content = message.reasoning_content %}
|
||||
{%- else %}
|
||||
{%- if think_end_token in content %}
|
||||
{%- set reasoning_content = content.split(think_end_token)[0].strip('\n').split(think_begin_token)[-1].strip('\n') %}
|
||||
{%- set content = content.split(think_end_token)[-1].strip('\n') %}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
|
||||
{%- if reasoning_content -%}
|
||||
{#- Render thinking for every assistant turn (all-turn visible) -#}
|
||||
{{- think_begin_token ~ reasoning_content ~ think_end_token }}
|
||||
{%- else -%}
|
||||
{#- No thinking rendered → prefix with think_end_token -#}
|
||||
{{- think_end_token }}
|
||||
{%- endif -%}
|
||||
|
||||
{%- if content -%}
|
||||
{{- content }}
|
||||
{%- endif -%}
|
||||
{%- if message.tool_calls -%}
|
||||
{{- toolcall_begin_token ~ '\n' }}
|
||||
|
||||
{%- for tool_call in message.tool_calls -%}
|
||||
{%- if tool_call.function -%}
|
||||
{%- set tool_call = tool_call.function -%}
|
||||
{%- endif -%}
|
||||
{{- ns_token + '<invoke name="' + tool_call.name + '">' }}
|
||||
{%- set _args = tool_call.arguments -%}
|
||||
{%- for k, v in _args.items() if v is not none %}
|
||||
{{- ns_token + '<' + k + '>' -}}
|
||||
{{- to_xml(v, ns_token) -}}
|
||||
{{- ns_token + '</' + k + '>' }}
|
||||
{%- endfor -%}
|
||||
{{- ns_token + '</invoke>' ~ '\n' }}
|
||||
{%- endfor -%}
|
||||
|
||||
{{- toolcall_end_token }}
|
||||
{%- if message.tool_calls[-1].function -%}
|
||||
{%- set last_tool_call.name = message.tool_calls[-1].function.name -%}
|
||||
{%- else -%}
|
||||
{%- set last_tool_call.name = message.tool_calls[-1].name -%}
|
||||
{%- endif -%}
|
||||
{%- else -%}
|
||||
{%- set last_tool_call.name = none -%}
|
||||
{%- endif -%}
|
||||
{{- eos_token ~ '\n' }}
|
||||
|
||||
{%- elif message.role == 'tool' -%}
|
||||
{%- if last_tool_call.name is none -%}
|
||||
{{- raise_exception("Message has tool role, but there was no previous assistant message with a tool call!") }}
|
||||
{%- endif -%}
|
||||
{%- if loop.first or (conversation_messages[loop.index0 - 1].role != 'tool') -%}
|
||||
{{- bos_token ~ 'tool' }}
|
||||
{%- endif -%}
|
||||
{{- '\n<response>' }}
|
||||
{%- if message.content is string -%}
|
||||
{{- message.content }}
|
||||
{%- else -%}
|
||||
{%- for tr in message.content -%}
|
||||
{%- if tr is mapping and tr.type is defined and tr.type == 'image' -%}
|
||||
{{- image_token }}
|
||||
{%- elif tr is mapping and tr.type is defined and tr.type == 'video' -%}
|
||||
{{- video_token }}
|
||||
{%- else -%}
|
||||
{{- tr.output if tr.output is defined else (tr.text if tr.type == 'text' and tr.text is defined else tr) }}
|
||||
{%- endif -%}
|
||||
{%- endfor -%}
|
||||
{%- endif -%}
|
||||
{{- '</response>' }}
|
||||
{%- if loop.last or (conversation_messages[loop.index0 + 1].role != 'tool') -%}
|
||||
{{- eos_token ~ '\n' -}}
|
||||
{%- endif -%}
|
||||
|
||||
{%- elif message.role == 'user' -%}
|
||||
{{- bos_token ~ 'user' ~ '\n' }}
|
||||
{{- visible_text(message.content) }}
|
||||
{{- eos_token ~ '\n' }}
|
||||
{%- endif -%}
|
||||
{%- endfor -%}
|
||||
|
||||
{#- Generation prompt -#}
|
||||
{%- if add_generation_prompt -%}
|
||||
{{- bos_token ~ 'ai' ~ '\n' }}
|
||||
{%- if thinking_mode is defined and thinking_mode == "disabled" -%}
|
||||
{{- think_end_token }}
|
||||
{%- elif thinking_mode is defined and thinking_mode == "adaptive" -%}
|
||||
{#- adaptive: no prefix, let model decide -#}
|
||||
{%- elif thinking_mode is defined and thinking_mode == "enabled" -%}
|
||||
{#- enabled or not defined: default to think -#}
|
||||
{{- think_begin_token }}
|
||||
{%- else -%}
|
||||
{#- adaptive: no prefix, let model decide -#}
|
||||
{%- endif -%}
|
||||
{%- endif -%}
|
||||
@@ -5,7 +5,7 @@ import os
|
||||
import sys
|
||||
import subprocess
|
||||
|
||||
HTTPLIB_VERSION = "refs/tags/v0.50.1"
|
||||
HTTPLIB_VERSION = "refs/tags/v0.51.0"
|
||||
|
||||
vendor = {
|
||||
"https://github.com/nlohmann/json/releases/latest/download/json.hpp": "vendor/nlohmann/json.hpp",
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
---
|
||||
name: add-new-model
|
||||
description: Guided workflow for adding a new model architecture to llama.cpp. Use when the user wants to add/port a new model architecture.
|
||||
---
|
||||
|
||||
# Add a new model architecture to llama.cpp
|
||||
|
||||
This skill walks a contributor through adding a new model architecture. AI-generated code is permitted in this project, so you may write full implementations for the steps below rather than only pointing at patterns - but follow `AGENTS.md`'s AI usage policy throughout:
|
||||
|
||||
- The contributor is 100% responsible for every line, however it was produced. They must be able to explain and defend any part of it to a reviewer. Check in with them as you go (don't silently generate everything and hand over a finished diff) so they actually absorb what was written.
|
||||
- Before writing code, make sure the contributor owns the design choices for this architecture (which reference model to follow, how non-standard bits like RoPE variants or MoE routing should be handled) - AI accelerates a design the contributor has already made, it doesn't make the design for them.
|
||||
- Disclosure is mandatory: any AI-meaningful contribution must be disclosed per the PR template. Remind the contributor of this before they open the PR.
|
||||
- Never write the PR description, commit message, GitHub issue/discussion post, or reviewer replies - those must come from the contributor. If asked to commit on their behalf, use `Assisted-by:` (never `Co-authored-by:`) and only after explicit confirmation.
|
||||
- If the requested change looks large or introduces a new pattern not covered here, pause and tell the user this kind of change is likely to need prior discussion with maintainers before a PR.
|
||||
- Keep the PR self-contained. If the work would require a lot of unconventional changes outside the new model file(s) (e.g. touching shared graph-building code, the sampler, or core APIs in ways other models don't), STOP and tell the contributor to open a discussion/issue first - invasive or excessive changes get closed without full review.
|
||||
- Do not bundle unrelated work into this PR - see Step 4 and Step 5 below for the specifics on multimodal and chat-template/parsing work.
|
||||
- Never hack around RoPE with a custom sin/cos implementation. Several past PRs tried this and were closed. If the existing `ggml_rope_ext` (see Step 2's RoPE tips) genuinely cannot express what this model needs, the contributor should open an issue to discuss it with maintainers first - not send a PR with a custom RoPE implementation.
|
||||
|
||||
Before starting, read `CONTRIBUTING.md`, `AGENTS.md` and `docs/development/HOWTO-add-model.md` if they are not already in context. Also run `git log --oneline -- src/models` and look at at least 3 recent PRs that added a model (their merge commits/diffs) - this shows current convention more reliably than the docs, which can lag behind.
|
||||
|
||||
## Step 0 - Scope and dedup check
|
||||
|
||||
Ask the contributor:
|
||||
1. Which model (HF repo id or name)? Is it text-only or does it have a multimodal (vision/audio) encoder?
|
||||
2. Do they already have the HF `config.json`/weights available locally?
|
||||
3. Have they checked for an existing PR/issue on this model? Suggest `gh search issues "<model name>"` and `gh search prs "<model name>"` in the `ggml-org/llama.cpp` repo. If an existing PR covers it, the contributor should comment there and collaborate rather than open a duplicate (per CONTRIBUTING.md's AI Usage Policy).
|
||||
4. What existing supported architecture is this model closest to (e.g. "Llama-like with sliding window", "MoE like DBRX", "BERT-style encoder")?
|
||||
|
||||
If the contributor doesn't know the closest reference architecture, you may grep `conversion/*.py` and `src/models/*.cpp` for architectures with a similar config shape (layer count, head count, MoE expert count, norm placement) and suggest 1-2 candidates - but let the contributor confirm the choice rather than picking one yourself; this choice is a design decision they need to own.
|
||||
|
||||
Do not proceed to Step 1 until the contributor has answered these and named a reference architecture.
|
||||
|
||||
## Step 1 - Convert the model to GGUF
|
||||
|
||||
Follow HOWTO-add-model.md section 1 for the actual touch points (conversion class registration, `constants.py`, `tensor_mapping.py`, etc.) - don't re-derive them here, read them from the doc.
|
||||
|
||||
Skill-specific addition: for each touch point, show the contributor the equivalent code in the reference architecture they named in Step 0 before writing the new version, and check that they understand what's different about their model (e.g. non-standard tensor shapes, extra hparams) rather than just copying the pattern silently.
|
||||
|
||||
## Step 2 - Define the architecture in llama.cpp
|
||||
|
||||
Follow HOWTO-add-model.md section 2 for the actual touch points (`llm_arch` enum, `LLM_ARCH_NAMES`, hparam loading, RoPE type case, etc.), including its "Tips and tricks" section for `ggml_rope_ext` gotchas.
|
||||
|
||||
Skill-specific addition: never hack around RoPE with a custom sin/cos implementation - see the RoPE rule above.
|
||||
|
||||
## Step 3 - Build the GGML graph
|
||||
|
||||
Follow HOWTO-add-model.md section 3 for the actual touch points (`src/models/<name>.cpp` struct, `llama_model_mapping` registration, etc.).
|
||||
|
||||
Skill-specific addition: before writing `src/models/<name>.cpp`, read at least 10 other files under `src/models/` (pick a mix, not just the one reference architecture) to confirm the struct layout, naming, and style you're about to write actually matches current convention - the pattern drifts over time and the HOWTO doc can lag behind it.
|
||||
|
||||
## Step 4 - Optional: multimodal encoder
|
||||
|
||||
Only do this if the contributor flagged a vision/audio encoder in Step 0. Follow HOWTO-add-model.md section 4 and `docs/multimodal.md` for the actual touch points (`MmprojModel` subclass, `clip.cpp`, `mtmd.cpp`, encoder graph in `tools/mtmd/models`, etc.).
|
||||
|
||||
Skill-specific addition, and read this carefully: **whether the multimodal encoder can be bundled into the same PR as the base text-model support depends on how conventional the change is.** It's OK to bundle it if the encoder support is conventional - i.e. no new infra or logic is needed, it's just a new cgraph reusing existing preprocessing/projector machinery (e.g. siglip/pixtral/qwen with just a new projector). If it requires anything beyond that - a new preprocessor, non-standard projector logic, or changes to shared `libmtmd` infra/logic - STOP, tell the contributor this is non-conventional, and have them land the text model first with the encoder as a dedicated follow-up PR. Do not let this decision pass silently - call it out explicitly to the contributor before writing any `clip.cpp`/`mtmd.cpp` code.
|
||||
|
||||
## Step 5 - Optional: chat template / parsing support
|
||||
|
||||
Only do this if the model needs a new built-in chat template (`src/llama-chat.cpp`) or a new output parser (see `docs/development/parsing.md` and `docs/autoparser.md`). If either is needed beyond what a user-supplied Jinja template already covers, treat it as its own dedicated follow-up PR, not part of the base model-support PR - call this out explicitly to the contributor rather than silently bundling it in.
|
||||
|
||||
## Common pitfalls (from past PR reviews)
|
||||
|
||||
These recur often enough in review comments on past add-model PRs that they're worth checking proactively, not just waiting for a reviewer to catch them:
|
||||
|
||||
- Don't validate the same hparam/config assumption in both the Python conversion script and the C++ load path - pick one layer to own the check, duplicating it just adds maintenance surface.
|
||||
- Optional hparams that are genuinely absent from some configs (e.g. a shared-expert count) should be read with an explicit optional/fallback accessor, not assumed present.
|
||||
- Hparams that are actually load-bearing (the model produces wrong output or crashes without them, e.g. `sliding_window_pattern`, norm-eps) must hard-error if missing, not silently fall back to a default.
|
||||
- Don't bake a default chat template into the C++ binary - inject it into the GGUF at conversion time instead, since one `llm_arch` can be reused by multiple fine-tunes with different templates, and a baked-in C++ default fails silently for those.
|
||||
- Before writing a dedicated tool-call/output parser, check whether the existing autoparser already handles the template (`llama-debug-template-parser <jinja>` shows what it detects).
|
||||
- Marking a custom EOS/closing-tag token as `eot` at conversion time isn't always sufficient - in long/agentic generations a model can emit the closing sequence as literal text instead of the token, so generation never stops on EOG and raw text leaks past the parser. Verify this case, not just the token path.
|
||||
- If reusing or aliasing an existing pre-tokenizer for convenience, justify and test that choice explicitly - silent reuse is an easy source of subtle tokenizer bugs.
|
||||
- Watch for excessive graph splits caused by building per-layer view/index tensors inside the layer loop - hoist tensors that don't vary per layer out of the loop (relevant if you hit `GGML_SCHED_MAX_SPLIT_INPUTS`).
|
||||
- A custom KQ mask fed into flash attention must match FA's expected dtype - cast it to F16 before passing it to `build_attn_mha` when FA is enabled.
|
||||
- When padding a custom KV-cache size to an alignment (e.g. `GGML_PAD(..., 256)`), apply the padding after all other size adjustments, not before - otherwise later logic can un-align it again.
|
||||
- For non-standard cache/SWA (sliding-window-attention) semantics, override the dedicated hook (e.g. `llama_model_n_swa()`) rather than mutating hparams to fake the behavior - hparams may be read elsewhere for unrelated purposes.
|
||||
- Don't ship unfinished or unverified speculative-decoding (e.g. MTP) scaffolding in the base model PR - if it hasn't actually been confirmed to work, pull it out and land it as its own follow-up.
|
||||
- Conversion code should call into the base class's existing hparam logic (e.g. `super().set_gguf_parameters()`) rather than re-deriving it - large blocks of code that duplicate what `TextModel`/`MmprojModel` already provide will get flagged as redundant.
|
||||
- Do constant tensor modifications (e.g. `norm(1 + weight)`) and permutations/chunking at conversion time, not in the graph - see HOWTO-add-model.md's "Prefer conversion-time tensor modifications" tip (Gemma 3 folds its `1 +` into the weights, Qwen3-Next permutes in `modify_tensors`). Doing these at runtime in the graph is very likely to be rejected as over-complicated; if you genuinely can't do it at conversion time, open a discussion first explaining why rather than implementing it in the graph.
|
||||
- Exception: a plain `weight * scale` with a constant scale is usually better applied at inference time instead of being folded into the weight at conversion. The scale conceptually applies to the activation, not the weight, so folding it in can hurt numerical stability, and it shifts the weight's value range in a way that can make quantization worse.
|
||||
|
||||
## Validation checklist
|
||||
|
||||
Reference: `examples/model-conversion/README.md`.
|
||||
|
||||
1. Convert to GGUF, then inspect/run both the original and converted tensors.
|
||||
2. Run logits verification (original vs converted). If this model is a new version of an already-supported family, verify the *previous* version still passes logits verification first - numerical differences may be pre-existing, not caused by the new work. The tools to perform full logits validation are available in `examples/model-conversion`.
|
||||
3. Quantize (including QAT variants if relevant) and re-verify.
|
||||
4. Run perplexity evaluation (simple and full).
|
||||
5. Sanity-check across `tools/cli`, `tools/completion`, `tools/imatrix`, `tools/quantize`, and `tools/server`.
|
||||
6. CPU backend first; other backends (CUDA, Metal, ...) can be separate follow-up PRs per `CONTRIBUTING.md`.
|
||||
7. Re-review every changed file against the coding/naming guidelines in `AGENTS.md` (and `CONTRIBUTING.md`'s "Coding guidelines"/"Naming guidelines" sections) - this is a separate pass from functional testing and is just as important: no forced line-wrapping, no unicode punctuation, minimal/non-redundant comments, `snake_case` naming (`kebab-case` for file names), matching indentation/brace style, etc.
|
||||
|
||||
## Before opening a PR
|
||||
|
||||
- Run the `code-review` skill on the diff first - it catches the convention and scope issues reviewers flag most often, and it's recommended to do this locally before pushing the PR.
|
||||
- Confirm the contributor can explain every changed line to a reviewer and is prepared to be asked about any of it - this is required regardless of how much of the code was AI-generated.
|
||||
- Confirm they did a comprehensive manual review of the full diff, not just a skim.
|
||||
- Fill in the AI-disclosure section of `.github/pull_request_template.md` describing how AI was used (do not omit or understate this).
|
||||
- Do not write the PR description, commit message, GitHub issue/discussion text, or any reviewer replies yourself - the contributor writes these.
|
||||
@@ -0,0 +1,143 @@
|
||||
---
|
||||
name: code-review
|
||||
description: Review llama.cpp changes against project conventions and common reviewer pitfalls before a PR. Use when the user wants to review a diff, branch, or PR.
|
||||
---
|
||||
|
||||
# Review llama.cpp changes
|
||||
|
||||
This skill reviews changes against llama.cpp's conventions and the pitfalls that reviewers flag most often, so the contributor can fix them before a maintainer has to. It has two modes:
|
||||
|
||||
- **Self-review (default):** review the contributor's own local changes (uncommitted work, or a branch vs `master`) as a pre-PR pass. Ask which if it's ambiguous; default to `git diff master...HEAD` plus any uncommitted changes.
|
||||
- **Read-only review of a PR/file:** if the user points at a PR number or specific files (including code they didn't write), review those and report findings.
|
||||
|
||||
In both modes the output is **private review notes for the user to read and act on** - it is never something to post. This is a hard rule from `AGENTS.md`: an agent must NEVER write, or help write, a PR comment, a review comment, or a reply to a reviewer, by any means including `gh`. Do not offer to. If the user asks you to post the notes, refuse and point them at that rule. Present findings in the conversation only.
|
||||
|
||||
Before starting, read `AGENTS.md` and `CONTRIBUTING.md` if not already in context - the "Coding guidelines", "Naming guidelines", and AI usage sections are the baseline this review enforces. For a diff that adds a new model architecture, also read `docs/development/HOWTO-add-model.md` and consider the dedicated `add-new-model` skill.
|
||||
|
||||
## Step 0 - Scope the diff and pick the checklists
|
||||
|
||||
Identify what actually changed and which area checklists below apply. Run `git diff --stat` (or `gh pr view <n> --json files` for PR mode) and bucket the touched paths:
|
||||
|
||||
- `conversion/`, `gguf-py/`, `src/models/`, `src/llama-arch.*` -> **New model / architecture**
|
||||
- `ggml/` (any backend, op, or `ggml.h`) -> **ggml / backend**
|
||||
- `include/llama.h` and other public headers -> **Public API**
|
||||
- `tools/server/` -> **Server**
|
||||
- anything else, plus all of the above -> **General** (always runs)
|
||||
|
||||
Always run the **Scope and quick-reject gate**, the **Security review**, and the **General** checklist. Run each area checklist whose paths were touched. Additionally, if the diff introduces a new component, subsystem, or piece of infrastructure (a new file/class/module, a new abstraction, or hand-rolled machinery), run the **Approach and design** review. Tell the user which checklists you're running and why.
|
||||
|
||||
## Scope and quick-reject gate (always)
|
||||
|
||||
These are the patterns that get PRs closed without a full review. Check them first - a finding here is more important than any code nit, because it can mean the change shouldn't be a PR in its current form at all.
|
||||
|
||||
- Is there a prior issue/discussion for this? Features are supposed to start as an issue, not a PR (`CONTRIBUTING.md`). If this is a nontrivial feature with no linked issue, flag it and suggest opening one first.
|
||||
- Is it a duplicate of existing/in-flight work? Suggest `gh search prs` / `gh search issues` for the feature. Many closed PRs were duplicates of something already queued.
|
||||
- Is it self-contained and single-purpose? Multiple unrelated changes/optimizations bundled together get sent back to be split. Flag unrelated changes and suggest separate PRs.
|
||||
- Does it touch multiple ggml backends at once? Initial support should be CPU-only, other backends as follow-ups (`CONTRIBUTING.md`). Flag CUDA/Metal/Vulkan/etc. changes bundled into a feature's first PR.
|
||||
- Does it add a new `ggml_type` / quantization type? That carries a disproportionate maintenance burden and needs the full justification package (GGUF sample upload, perplexity vs FP16/BF16 and similar sizes, KL-divergence data, CPU perf numbers). Absent that, it will be rejected regardless of code quality.
|
||||
- Is it invasive - new subsystem, core-API reshaping, changes to shared graph/sampler code that other models don't need? Flag it and suggest a discussion with maintainers before investing further.
|
||||
- Is it niche/vendor-specific in a way that adds a maintenance burden nobody will own long-term? Flag the maintenance-ownership question.
|
||||
- Is the change semantically correct, or a plausible-looking "fix" that misunderstands the code? Sanity-check the actual behavior, not just that it compiles.
|
||||
- AI-disclosure: if AI meaningfully contributed, is the PR template's disclosure section filled in? Remind the user. Never suggest writing the PR description or commit message for them.
|
||||
|
||||
## Security review (mandatory)
|
||||
|
||||
Mandatory on every review; any finding here is **blocking**. Rule of thumb: GGUF metadata, tensor shapes, tokenizer/grammar input, and all server/RPC fields are attacker-controlled - bound them before use.
|
||||
|
||||
- **Sizes/counts from tensor dims:** validate before allocating. Products like `ne[i]*nb[i]`/nbytes can overflow on crafted dims into an undersized alloc then heap overflow. Overflow checks must run BEFORE the arithmetic they guard - padding/alignment macros wrap to 0 near `SIZE_MAX`, so a guard after the pad passes.
|
||||
- **GGUF strings/arrays:** cap declared lengths and element counts before using them to size a loop or buffer; validate element type and length before casting an array to a pointer or reading fixed indices (`[i+1]`, `[0..2]`).
|
||||
- **File-supplied counts indexing fixed arrays:** bound any count (e.g. layer/block count into a `LLAMA_MAX_*` array) before indexing; watch checks that only fire when an optional key is present.
|
||||
- **Bounds comparisons:** flag narrowing casts (`size_t`->`int32_t`) and signed/unsigned mixing that can bypass a length check and copy past a buffer.
|
||||
- **Parsed/derived indices:** range-check `stoi`/`atoi` results and catch parse throws; never use a default or derived token id (EOS/BOS/...) as an index without a bounds check.
|
||||
- **Reused/reserved buffers:** recheck bounds after a buffer is shrunk or reused; watch `reserve()` then index-by-assumed-size, and header fields read before their length is checked.
|
||||
- **Server JSON ints:** clamp client-supplied integers (token/discard counts, offsets) to non-negative and an upper bound before they reach index/pointer arithmetic.
|
||||
- **RPC-deserialized fields:** treat every field (type/buffer/data/ne/nb/op_params) as hostile - validate before use. Null/zero buffers skipping validation, attacker data pointers, out-of-range type indices, and negative strides sign-extending past a corner-only assert all give arbitrary read/write.
|
||||
- **Lifetime/UAF:** flag stored raw pointers to caller/temporary storage, cached pointers to buffers a later free releases, async ops whose source may drop before completion, and structures not invalidated on free/realloc. Null-check conditionally-built or "not required" tensors before dereferencing.
|
||||
|
||||
## Approach and design (when a new component/infra is introduced)
|
||||
|
||||
Run this whenever the diff adds a new component, subsystem, or piece of infrastructure. Reviews too often stop at "does it work" - a diff can be correct and still be the wrong approach, and a messy design costs more long-term than a bug. Evaluate the *approach*, not just the behavior; raising a cleaner one is a high-value finding, not a nit. If you see a better design, describe it concretely rather than just calling the current one bad.
|
||||
|
||||
- **Simpler approach upstream:** the biggest win is often a different data model or design that removes whole subsystems, not tweaks to the code as written. Complexity must be justified by the problem, not by the first thing that worked.
|
||||
- **Reuse over reinvention:** grep for an existing helper, library, object, or mechanism before adding a new one. Reimplementing what the codebase already has reintroduces solved bugs and adds maintenance surface.
|
||||
- **Clear ownership/lifetime:** prefer RAII and obvious ownership over manual liveness flags, hand-tracked pointers, and "is it still alive?" checks - manual lifetime tracking is a recurring source of subtle bugs.
|
||||
- **Right-sized machinery:** flag redundant, overkill, or heavier-than-needed primitives and abstractions; use the minimum the design actually needs.
|
||||
- **Right structure and fit:** a new type should earn its place (split it if it serves two roles); follow existing patterns, idioms, and naming, and avoid constructs the project shuns.
|
||||
- **Root cause vs symptom:** fixes layered on fixes signal a design to correct, not guard around.
|
||||
|
||||
## New model / architecture
|
||||
|
||||
See the `add-new-model` skill and `docs/development/HOWTO-add-model.md` for the full workflow; this is the review-time subset that reviewers most often catch:
|
||||
|
||||
- Don't branch on `model.arch` when the real dependency is a config/capability value - gate on the hparam/capability, not the architecture enum.
|
||||
- If the model is a close variant of an existing arch, is the delta justified? Prefer reusing or subclassing the existing arch/model class over duplicating it. A near-duplicate class or `src/models/<name>.cpp` will be asked to merge with its sibling.
|
||||
- New tensor names go through `tensor_mapping.py`, not ad-hoc name matching.
|
||||
- For QKV, split the *activation* with `ggml_view`, not the *weight* tensor; rely on ggml broadcasting instead of manually duplicating tensors.
|
||||
- New graph inputs are declared at the top of the graph-build function, not inline where first used.
|
||||
- Hparams that the model can't run correctly without must be mandatory (hard-error if missing), not read with a silent default fallback. Only genuinely-optional-across-configs values get a fallback accessor.
|
||||
- New/optional weight tensors (scales, etc.) must route through `build_lora_mm` and the existing helpers, matching convention - don't leave raw matmuls copied from another arch.
|
||||
- Don't hack RoPE with a custom sin/cos implementation. If `ggml_rope_ext` genuinely can't express it, that's an issue for discussion, not a PR.
|
||||
- Test the quantized-KV path (`-ctk`/`-ctv q8_0`), not just default f16 - new speculative/attention features silently break there.
|
||||
- Preserve existing explanatory comments about model-specific quirks when copying code; note the provenance ("copied from X, with Y added").
|
||||
- Remove dead code/branches left over from adapting a reference implementation.
|
||||
|
||||
## ggml / backend
|
||||
|
||||
- `supports_op` (and any dispatch/gating condition) must be scoped exactly to the cases being changed - a condition meant for a few quant types must not silently disable or enable everything else.
|
||||
- No hardcoded warp/lane size - use `ggml_cuda_get_physical_warp_size()` (32 on CUDA, 64 on HIP/ROCm) and the portable helpers.
|
||||
- Strip leftover debug/profiling/logging code before review.
|
||||
- New or changed op? Update `docs/ops.md` and the relevant `docs/ops/*.csv` for the touched backend.
|
||||
- New op or operator change needs corresponding `test-backend-ops` cases, and (per `CONTRIBUTING.md`) consistency across at least two backends.
|
||||
- New kernels are expected to come with concrete perf data (throughput across realistic tensor shapes), not just correctness.
|
||||
- Don't have a backend mutate the cgraph as a shortcut - that's an unresolved architectural question, not something to slip in.
|
||||
- Expect this to need two maintainer approvals; that's normal for `ggml/` changes, not a sign something is wrong.
|
||||
- For CUDA: Avoid excessively templating kernels, only add this where it shows visible performance gain.
|
||||
|
||||
## Public API (`include/llama.h`)
|
||||
|
||||
Public API changes carry a higher bar than internal ones (`CONTRIBUTING.md`). Review for:
|
||||
|
||||
- Justification: why doesn't an existing mechanism (e.g. `cb_eval`, existing batch/sampler knobs) suffice? If it does, the change likely shouldn't add public surface. This is the single most common reason these PRs are rejected.
|
||||
- Experimental or stop-gap surface belongs in a side header (`llama-ext.h`), not in `llama.h`.
|
||||
- Keep it minimal and general: prefer one general call over several narrow convenience wrappers; make new calls forward-compatible (e.g. mixed-modality batches) rather than assuming today's shape.
|
||||
- The C API is the first-class, stable, ABI-defining surface - don't propose a parallel C++ API as a replacement. `llama-cpp.h` stays a thin convenience layer.
|
||||
- Types and naming: sized integer types (`int32_t`, `size_t` for sizes/offsets); `snake_case`; `<class>_<method>` = `<class>_<action>_<noun>`; enum values upper-case and prefixed with the enum name; `_t` suffix for opaque types. Avoid gratuitous signature/ABI changes to existing exported functions.
|
||||
- Every new API needs a working example/tool exercising it in the same PR - reviewers find real bugs by requiring it to be wired into `server`, `embedding`, `perplexity`, etc.
|
||||
|
||||
## Server (`tools/server/`)
|
||||
|
||||
- Is the feature within server's defined scope? Check `tools/server/README-dev.md` - out-of-scope features get declined.
|
||||
- Security: don't trust client-supplied headers (e.g. `X-Forwarded-For`) or add footguns; things like IP allowlisting belong at a reverse proxy unless there's a trusted-proxy design.
|
||||
- Wire new behavior into the existing request/response and checkpoint paths correctly; watch for resource leaks across requests.
|
||||
|
||||
## Multimodal (`tools/mtmd/`)
|
||||
|
||||
- Tensor names must be prefixed by `v.`, `a.`, `mm.` or `a.mm.` (legacy naming doesn't follow this convention - this is expected, but new code should follow it).
|
||||
- Do not use explicit sin/cos for RoPE; use `ggml_rope_ext` instead, see `HOWTO-add-model.md`. If it can't express the needed behavior, that's a design discussion, not a PR.
|
||||
- New GGML ops must not be introduced in the same PR, you must push it as a separate PR.
|
||||
- In most cases, `build_vit` should be enough to build the transformer graph for vision models. Do not add a loop to build the transformer graph manually, unless you have a very good reason to do so. If you do, please explain why in the PR description.
|
||||
- If you need a dedicated preprocessor, there is a high chance that it can be a derived class from one of the existing preprocessors. Check carefully before adding a new preprocessor class.
|
||||
- If the model need a new public API in `mtmd.h`, open a discussion first.
|
||||
|
||||
## General (always)
|
||||
|
||||
Enforce the `AGENTS.md` / `CONTRIBUTING.md` coding and naming guidelines on every changed line - this is a distinct pass from checking that the code works, and matters just as much for review speed:
|
||||
|
||||
- ASCII only in code and comments - no emdash, unicode arrows, `x`, `...` used as unicode; use `-`, `->`, `x`, `...` ASCII equivalents.
|
||||
- Comments are concise and explain non-obvious *why*, not *what*. Flag verbose comments, comments that restate the code, comments that reference the current task/PR, and comments hard-wrapped to a fixed column width.
|
||||
- Do not force-wrap prose/comments to a fixed character count or split a sentence across lines.
|
||||
- `snake_case` names; `kebab-case` (lowercase-with-dashes) file names for C/C++, `.h` headers; Python files lowercase-with-underscores. Naming optimizes for longest common prefix (`number_small`, not `small_number`).
|
||||
- 4-space indentation, brackets on the same line, `void * ptr`, `int & a`, no trailing whitespace; match the surrounding style.
|
||||
- Reuse existing infrastructure over introducing new components; no new third-party dependencies, extra headers, or files unless clearly justified.
|
||||
- Keep it simple: a simpler change doing 90% is often preferable to a complex one doing 100%. Flag unnecessary templates/fancy STL; basic `for` loops are fine here.
|
||||
- Every added line should be something the contributor can explain and defend to a reviewer without AI help - flag anything that looks copied-in without understanding.
|
||||
|
||||
## Reporting
|
||||
|
||||
Group findings by severity so the user knows what actually blocks a merge:
|
||||
|
||||
1. **Blocking** - quick-reject/scope issues and correctness bugs; these can sink the PR regardless of everything else.
|
||||
2. **Will slow the review** - convention/naming/comment violations, missing tests/docs/perf data, missing API justification or example.
|
||||
3. **Nits** - minor style, optional cleanups.
|
||||
|
||||
For each finding, point to the file and line and say concretely what to change and why. Do not rewrite the whole diff unprompted; let the contributor make the fixes so they own and understand them. And do not draft any PR text, commit message, or reviewer reply - that is the contributor's to write.
|
||||
@@ -127,6 +127,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
|
||||
{ LLM_ARCH_GROVEMOE, "grovemoe" },
|
||||
{ LLM_ARCH_APERTUS, "apertus" },
|
||||
{ LLM_ARCH_MINIMAX_M2, "minimax-m2" },
|
||||
{ LLM_ARCH_MINIMAX_M3, "minimax-m3" },
|
||||
{ LLM_ARCH_COGVLM, "cogvlm" },
|
||||
{ LLM_ARCH_RND1, "rnd1" },
|
||||
{ LLM_ARCH_PANGU_EMBED, "pangu-embedded" },
|
||||
@@ -142,6 +143,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
|
||||
{ LLM_ARCH_KIMI_LINEAR, "kimi-linear" },
|
||||
{ LLM_ARCH_TALKIE, "talkie" },
|
||||
{ LLM_ARCH_MELLUM, "mellum" },
|
||||
{ LLM_ARCH_NANBEIGE, "nanbeige" },
|
||||
{ LLM_ARCH_UNKNOWN, "(unknown)" },
|
||||
};
|
||||
|
||||
@@ -220,6 +222,8 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" },
|
||||
{ LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" },
|
||||
{ LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" },
|
||||
{ LLM_KV_NUM_LOOPS, "%s.num_loops" },
|
||||
{ LLM_KV_SKIP_LOOP_FINAL_NORM, "%s.skip_loop_final_norm" },
|
||||
|
||||
{ LLM_KV_ATTENTION_HEAD_COUNT, "%s.attention.head_count" },
|
||||
{ LLM_KV_ATTENTION_HEAD_COUNT_KV, "%s.attention.head_count_kv" },
|
||||
@@ -253,6 +257,9 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" },
|
||||
{ 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" },
|
||||
@@ -310,6 +317,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_TARGET_LAYERS, "%s.target_layers" },
|
||||
{ LLM_KV_TARGET_HIDDEN_SIZE, "%s.target_hidden_size" },
|
||||
{ LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" },
|
||||
{ LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" },
|
||||
|
||||
{ LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" },
|
||||
// sentence-transformers dense modules feature dims
|
||||
@@ -596,6 +604,9 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
|
||||
{ LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" },
|
||||
{ LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" },
|
||||
{ LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" },
|
||||
{ LLM_TENSOR_INDEXER_Q_PROJ, "blk.%d.indexer.q_proj" },
|
||||
{ LLM_TENSOR_INDEXER_K_PROJ, "blk.%d.indexer.k_proj" },
|
||||
{ LLM_TENSOR_INDEXER_Q_NORM, "blk.%d.indexer.q_norm" },
|
||||
{ LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "blk.%d.indexer_compressor_kv" },
|
||||
{ LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "blk.%d.indexer_compressor_gate" },
|
||||
{ LLM_TENSOR_INDEXER_COMPRESSOR_APE, "blk.%d.indexer_compressor_ape" },
|
||||
@@ -605,6 +616,9 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
|
||||
{ LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" },
|
||||
{ LLM_TENSOR_FC, "fc" },
|
||||
{ LLM_TENSOR_D2T, "d2t" },
|
||||
{ LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" },
|
||||
{ LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" },
|
||||
{ LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" },
|
||||
};
|
||||
|
||||
// declare information about the model weight tensors:
|
||||
@@ -831,6 +845,9 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
|
||||
{LLM_TENSOR_INDEXER_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_Q_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_K_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
|
||||
{LLM_TENSOR_INDEXER_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
|
||||
@@ -856,6 +873,10 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
|
||||
// eagle3
|
||||
{LLM_TENSOR_FC, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_D2T, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
|
||||
// dspark
|
||||
{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_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {}
|
||||
@@ -1000,6 +1021,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
|
||||
case LLM_ARCH_LFM2:
|
||||
case LLM_ARCH_LFM2MOE:
|
||||
case LLM_ARCH_MINIMAX_M2:
|
||||
case LLM_ARCH_MINIMAX_M3:
|
||||
case LLM_ARCH_MISTRAL4:
|
||||
case LLM_ARCH_KIMI_LINEAR:
|
||||
return false;
|
||||
|
||||
@@ -146,7 +146,9 @@ enum llm_arch {
|
||||
LLM_ARCH_TALKIE,
|
||||
LLM_ARCH_MELLUM,
|
||||
LLM_ARCH_EAGLE3,
|
||||
LLM_ARCH_MINIMAX_M3,
|
||||
LLM_ARCH_DFLASH,
|
||||
LLM_ARCH_NANBEIGE,
|
||||
LLM_ARCH_UNKNOWN,
|
||||
};
|
||||
|
||||
@@ -225,6 +227,8 @@ enum llm_kv {
|
||||
LLM_KV_TOKEN_SHIFT_COUNT,
|
||||
LLM_KV_INTERLEAVE_MOE_LAYER_STEP,
|
||||
LLM_KV_FULL_ATTENTION_INTERVAL,
|
||||
LLM_KV_NUM_LOOPS,
|
||||
LLM_KV_SKIP_LOOP_FINAL_NORM,
|
||||
|
||||
LLM_KV_ATTENTION_HEAD_COUNT,
|
||||
LLM_KV_ATTENTION_HEAD_COUNT_KV,
|
||||
@@ -258,6 +262,9 @@ enum llm_kv {
|
||||
LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,
|
||||
LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,
|
||||
LLM_KV_ATTENTION_INDEXER_TOP_K,
|
||||
LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE,
|
||||
LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS,
|
||||
LLM_KV_ATTENTION_INDEXER_TYPES,
|
||||
LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,
|
||||
LLM_KV_ATTENTION_OUTPUT_LORA_RANK,
|
||||
LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,
|
||||
@@ -356,6 +363,7 @@ enum llm_kv {
|
||||
LLM_KV_TARGET_LAYERS,
|
||||
LLM_KV_TARGET_HIDDEN_SIZE,
|
||||
LLM_KV_NORM_BEFORE_RESIDUAL,
|
||||
LLM_KV_NORM_BEFORE_FC,
|
||||
|
||||
LLM_KV_SHORTCONV_L_CACHE,
|
||||
|
||||
@@ -596,6 +604,9 @@ enum llm_tensor {
|
||||
LLM_TENSOR_INDEXER_PROJ,
|
||||
LLM_TENSOR_INDEXER_ATTN_K,
|
||||
LLM_TENSOR_INDEXER_ATTN_Q_B,
|
||||
LLM_TENSOR_INDEXER_Q_PROJ,
|
||||
LLM_TENSOR_INDEXER_K_PROJ,
|
||||
LLM_TENSOR_INDEXER_Q_NORM,
|
||||
LLM_TENSOR_INDEXER_COMPRESSOR_WKV,
|
||||
LLM_TENSOR_INDEXER_COMPRESSOR_WGATE,
|
||||
LLM_TENSOR_INDEXER_COMPRESSOR_APE,
|
||||
@@ -613,6 +624,9 @@ enum llm_tensor {
|
||||
LLM_TENSOR_MASKED_EMBD_ORDERING,
|
||||
LLM_TENSOR_FC,
|
||||
LLM_TENSOR_D2T,
|
||||
LLM_TENSOR_DSPARK_MARKOV_W1,
|
||||
LLM_TENSOR_DSPARK_MARKOV_W2,
|
||||
LLM_TENSOR_DSPARK_CONF_PROJ,
|
||||
};
|
||||
|
||||
|
||||
|
||||
@@ -2338,7 +2338,9 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
|
||||
model.arch == LLM_ARCH_KIMI_LINEAR ||
|
||||
model.arch == LLM_ARCH_QWEN35 ||
|
||||
model.arch == LLM_ARCH_QWEN35MOE ||
|
||||
model.arch == LLM_ARCH_DEEPSEEK4) {
|
||||
model.arch == LLM_ARCH_DEEPSEEK4 ||
|
||||
model.arch == LLM_ARCH_NANBEIGE ||
|
||||
model.arch == LLM_ARCH_MINIMAX_M3) {
|
||||
return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
|
||||
}
|
||||
uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
|
||||
@@ -2472,11 +2474,12 @@ llm_graph_cb llama_context::graph_get_cb() const {
|
||||
ggml_set_name(cur, name);
|
||||
}
|
||||
|
||||
// norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends
|
||||
// - norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends
|
||||
// - force the last op of the layer on the specified backend to avoid running it on the backend of the next layer due to scheduling
|
||||
// FIXME: fix in ggml_backend_sched
|
||||
const bool full_offload = model.n_gpu_layers() > model.hparams.n_layer_all;
|
||||
if (ubatch.n_tokens < 32 || full_offload) {
|
||||
if (il != -1 && strcmp(name, "norm") == 0) {
|
||||
if (il != -1 && (strcmp(name, "norm") == 0 || strcmp(name, "l_last") == 0)) {
|
||||
const auto & dev_layer = model.dev_layer(il);
|
||||
for (const auto & backend : backends) {
|
||||
if (ggml_backend_get_device(backend.get()) == dev_layer) {
|
||||
|
||||
@@ -1139,6 +1139,18 @@ struct llama_grammar * llama_grammar_init_impl(
|
||||
vec_rules[i].push_back({LLAMA_GRETYPE_END, 0});
|
||||
}
|
||||
|
||||
// Validate that all rule references point to valid rules
|
||||
for (size_t i = 0; i < n_rules; i++) {
|
||||
for (const auto & elem : vec_rules[i]) {
|
||||
if (elem.type == LLAMA_GRETYPE_RULE_REF) {
|
||||
if (elem.value >= n_rules || vec_rules[elem.value].empty()) {
|
||||
LLAMA_LOG_ERROR("invalid grammar: rule %zu references undefined rule %u\n", i, elem.value);
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check for left recursion
|
||||
std::vector<bool> rules_visited(n_rules);
|
||||
std::vector<bool> rules_in_progress(n_rules);
|
||||
|
||||
+12
-1
@@ -1709,6 +1709,17 @@ ggml_tensor * llm_graph_context::build_ffn(
|
||||
cur = ggml_swiglu(ctx0, cur);
|
||||
cb(cur, "ffn_swiglu", il);
|
||||
} break;
|
||||
case LLM_FFN_SWIGLU_OAI_MOE:
|
||||
if (gate && type_gate == LLM_FFN_PAR) {
|
||||
// same alpha/limit constants as gpt-oss
|
||||
const float alpha = 1.702f;
|
||||
const float limit = 7.0f;
|
||||
cur = ggml_swiglu_oai(ctx0, cur, tmp, alpha, limit);
|
||||
cb(cur, "ffn_swiglu_oai", il);
|
||||
type_gate = LLM_FFN_SEQ;
|
||||
} else {
|
||||
GGML_ABORT("LLM_FFN_SWIGLU_OAI_MOE requires a parallel gate");
|
||||
} break;
|
||||
case LLM_FFN_GEGLU:
|
||||
{
|
||||
cur = ggml_geglu(ctx0, cur);
|
||||
@@ -2668,7 +2679,7 @@ ggml_tensor * llm_graph_context::build_attn(
|
||||
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
|
||||
}
|
||||
|
||||
const auto & kq_mask = inp->get_kq_mask();
|
||||
ggml_tensor * kq_mask = inp->get_kq_mask();
|
||||
|
||||
ggml_tensor * q = q_cur;
|
||||
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
|
||||
|
||||
@@ -180,6 +180,16 @@ uint32_t llama_hparams::n_embd_v_gqa_max() const {
|
||||
return val;
|
||||
}
|
||||
|
||||
uint32_t llama_hparams::n_embd_k_idx(uint32_t il) const {
|
||||
if (!indexer_kv || indexer_head_size == 0) {
|
||||
return 0; // arch without a MSA indexer
|
||||
}
|
||||
if (il < n_layer_dense_lead) {
|
||||
return 0; // leading dense layers carry no indexer
|
||||
}
|
||||
return indexer_head_size; // 128
|
||||
}
|
||||
|
||||
uint32_t llama_hparams::n_embd_r() const {
|
||||
if (wkv_head_size != 0) {
|
||||
// for RWKV models
|
||||
@@ -248,6 +258,14 @@ bool llama_hparams::is_mla() const {
|
||||
return n_embd_head_k_mla_impl != 0 && n_embd_head_v_mla_impl != 0;
|
||||
}
|
||||
|
||||
bool llama_hparams::is_indexer_full(uint32_t il) const {
|
||||
if (il < n_layer()) {
|
||||
return is_indexer_full_impl[il];
|
||||
}
|
||||
|
||||
GGML_ABORT("%s: il (%u) out of bounds (n_layer: %u)\n", __func__, il, n_layer());
|
||||
}
|
||||
|
||||
uint32_t llama_hparams::n_embd_head_k_mla() const {
|
||||
return is_mla() ? n_embd_head_k_mla_impl : n_embd_head_k();
|
||||
}
|
||||
|
||||
@@ -47,6 +47,7 @@ struct llama_hparams {
|
||||
bool use_par_res;
|
||||
bool swin_norm;
|
||||
bool norm_before_residual = false;
|
||||
bool norm_before_fc = false;
|
||||
|
||||
uint32_t n_ctx_train; // context size the model was trained on
|
||||
uint32_t n_embd;
|
||||
@@ -226,6 +227,15 @@ struct llama_hparams {
|
||||
uint32_t indexer_n_head = 0;
|
||||
uint32_t indexer_head_size = 0;
|
||||
uint32_t indexer_top_k = 0;
|
||||
// MSA
|
||||
uint32_t indexer_block_size = 0;
|
||||
uint32_t indexer_local_blocks = 0;
|
||||
// MSA stores its indexer keys in the main KV cache (k_idx tensors);
|
||||
bool indexer_kv = false;
|
||||
|
||||
// Indexer is "full" (1) or "shared" (0)
|
||||
// Shared indexers reuse top-k from previous full layer
|
||||
std::array<uint32_t, LLAMA_MAX_LAYERS> is_indexer_full_impl;
|
||||
|
||||
// DeepSeek-V4
|
||||
uint32_t dsv4_o_group_count = 0;
|
||||
@@ -302,6 +312,8 @@ struct llama_hparams {
|
||||
|
||||
bool is_swa(uint32_t il) const;
|
||||
|
||||
bool is_indexer_full(uint32_t il) const;
|
||||
|
||||
void set_recr_pattern(uint32_t n_pattern, bool dense_first = false);
|
||||
|
||||
// whether or not the given layer is recurrent (for hybrid models)
|
||||
@@ -344,6 +356,9 @@ struct llama_hparams {
|
||||
uint32_t n_embd_k_gqa_max() const;
|
||||
uint32_t n_embd_v_gqa_max() const;
|
||||
|
||||
// dimension of the single-head MSA indexer key stream
|
||||
uint32_t n_embd_k_idx(uint32_t il = 0) const;
|
||||
|
||||
// dimension of the rolling state embeddings
|
||||
// corresponds to Mamba's conv_states size or RWKV's token_shift states size
|
||||
uint32_t n_embd_r() const;
|
||||
|
||||
+285
-16
@@ -112,7 +112,7 @@ llama_kv_cache::llama_kv_cache(
|
||||
auto it = ctx_map.find(buft);
|
||||
if (it == ctx_map.end()) {
|
||||
ggml_init_params params = {
|
||||
/*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()),
|
||||
/*.mem_size =*/ size_t(3u*(1 + n_stream)*n_layer*ggml_tensor_overhead()), //Reserve tensor metadata for up to 3 tensors per layer (K, V, and optional K_idx), plus one view per tensor per stream.
|
||||
/*.mem_buffer =*/ NULL,
|
||||
/*.no_alloc =*/ true,
|
||||
};
|
||||
@@ -242,9 +242,25 @@ llama_kv_cache::llama_kv_cache(
|
||||
v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr);
|
||||
}
|
||||
|
||||
const uint32_t n_embd_k_idx = hparams.n_embd_k_idx(il);
|
||||
ggml_tensor * k_idx = n_embd_k_idx > 0
|
||||
? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_k_idx, kv_size, n_stream)
|
||||
: nullptr;
|
||||
if (k_idx) {
|
||||
ggml_format_name(k_idx, "cache_k_idx_l%d", il);
|
||||
msa_strict_slots = (n_stream == n_seq_max);
|
||||
}
|
||||
|
||||
std::vector<ggml_tensor *> k_idx_stream;
|
||||
for (uint32_t s = 0; s < n_stream; ++s) {
|
||||
k_idx_stream.push_back(k_idx
|
||||
? ggml_view_2d(ctx, k_idx, n_embd_k_idx, kv_size, k_idx->nb[1], s*k_idx->nb[2])
|
||||
: nullptr);
|
||||
}
|
||||
|
||||
map_layer_ids[il] = layers.size();
|
||||
|
||||
layers.push_back({ il, k, v, k_stream, v_stream, });
|
||||
layers.push_back({ il, k, v, k_idx, k_stream, v_stream, k_idx_stream });
|
||||
}
|
||||
|
||||
if (reuse) {
|
||||
@@ -293,13 +309,24 @@ llama_kv_cache::llama_kv_cache(
|
||||
}
|
||||
|
||||
{
|
||||
const size_t memory_size_k = size_k_bytes();
|
||||
const size_t memory_size_v = size_v_bytes();
|
||||
const size_t memory_size_k = size_k_bytes();
|
||||
const size_t memory_size_v = size_v_bytes();
|
||||
const size_t memory_size_k_idx = size_k_idx_bytes();
|
||||
const size_t memory_size_total = memory_size_k + memory_size_v + memory_size_k_idx;
|
||||
|
||||
LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__,
|
||||
(float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream,
|
||||
ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f),
|
||||
ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f));
|
||||
constexpr float mib = 1024.0f * 1024.0f;
|
||||
|
||||
const std::string k_log = format(", K (%s): %7.2f MiB", ggml_type_name(type_k), (float) memory_size_k / mib);
|
||||
const std::string v_log = format(", V (%s): %7.2f MiB", ggml_type_name(type_v), (float) memory_size_v / mib);
|
||||
|
||||
std::string k_idx_log;
|
||||
if (memory_size_k_idx > 0) {
|
||||
k_idx_log = format(", K_idx (%s): %7.2f MiB", ggml_type_name(GGML_TYPE_F32), (float) memory_size_k_idx / mib);
|
||||
}
|
||||
|
||||
LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs)%s%s%s\n", __func__,
|
||||
(float) memory_size_total / mib, kv_size, (int) layers.size(), n_seq_max, n_stream,
|
||||
k_log.c_str(), v_log.c_str(), k_idx_log.c_str());
|
||||
}
|
||||
|
||||
// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
|
||||
@@ -323,7 +350,7 @@ llama_kv_cache::llama_kv_cache(
|
||||
hparams.n_embd_head_k() % 64 == 0;
|
||||
|
||||
// always create Hadamard rotation tensors for DeepSeek lightning indexers
|
||||
if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4) &&
|
||||
if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4 || model.arch == LLM_ARCH_GLM_DSA) &&
|
||||
hparams.n_embd_head_k_full == hparams.indexer_head_size) {
|
||||
attn_rot_k = true;
|
||||
}
|
||||
@@ -392,6 +419,39 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
|
||||
p1 = std::numeric_limits<llama_pos>::max();
|
||||
}
|
||||
|
||||
// empty range - nothing to remove
|
||||
if (p0 >= p1) {
|
||||
return true;
|
||||
}
|
||||
|
||||
// MSA anchors block selection to absolute cache slots (slot == position). Tail trim and full removal preserve this invariant, but removing a prefix
|
||||
// or middle range would free slots while later cells survive, desynchronizing the indexer cache. Reject such removals before modifying the cache.
|
||||
if (msa_strict_slots) {
|
||||
for (llama_seq_id sid = 0; sid < (llama_seq_id) seq_to_stream.size(); ++sid) {
|
||||
if (seq_id >= 0 && sid != seq_id) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const auto & cells = v_cells[seq_to_stream[sid]];
|
||||
|
||||
const llama_pos pmin = cells.seq_pos_min(sid);
|
||||
const llama_pos pmax = cells.seq_pos_max(sid);
|
||||
|
||||
if (pmin < 0) {
|
||||
continue; // empty sequence
|
||||
}
|
||||
|
||||
const bool overlaps = p0 <= pmax && p1 > pmin; // the range removes something
|
||||
const bool leaves_tail = p1 <= pmax; // cells beyond the range survive
|
||||
|
||||
if (overlaps && leaves_tail) {
|
||||
LLAMA_LOG_WARN("%s: MSA: partial (non-suffix) removal [%d, %d) for seq %d is not supported "
|
||||
"(block selection is anchored to cache slots) - rejected\n", __func__, p0, p1, sid);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (seq_id >= 0) {
|
||||
auto & cells = v_cells[seq_to_stream[seq_id]];
|
||||
auto & head = v_heads[seq_to_stream[seq_id]];
|
||||
@@ -846,6 +906,10 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co
|
||||
if (layer.v_stream[ssrc]) {
|
||||
ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]);
|
||||
}
|
||||
if (layer.k_idx_stream[ssrc]) {
|
||||
GGML_ASSERT(layer.k_idx_stream[sdst]);
|
||||
ggml_backend_tensor_copy(layer.k_idx_stream[ssrc], layer.k_idx_stream[sdst]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -994,6 +1058,44 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
|
||||
|
||||
const auto & cells = v_cells[seq_to_stream[seq_id]];
|
||||
|
||||
if (n_tokens > cells.size()) {
|
||||
LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());
|
||||
return { };
|
||||
}
|
||||
|
||||
// MSA block selection assumes slot == logical position (append-only streams).
|
||||
if (msa_strict_slots) {
|
||||
for (uint32_t ii = 0; ii < n_tokens; ++ii) {
|
||||
const llama_pos pos = ubatch.pos[s*n_tokens + ii];
|
||||
|
||||
if (pos < 0 || (uint64_t) pos >= cells.size()) {
|
||||
LLAMA_LOG_WARN("%s: MSA: position %d is outside the cache range [0, %u)\n",
|
||||
__func__, pos, cells.size());
|
||||
return { };
|
||||
}
|
||||
|
||||
const uint32_t idx = (uint32_t) pos;
|
||||
|
||||
if (!cells.is_empty(idx)) {
|
||||
LLAMA_LOG_WARN("%s: MSA: required slot %u is already occupied (stream %u)\n",
|
||||
__func__, idx, seq_to_stream[seq_id]);
|
||||
return { };
|
||||
}
|
||||
|
||||
// strictly increasing positions, rules out duplicates and, for contiguous requests, is tightened to exact adjacency
|
||||
if (!res.idxs[s].empty() && (cont ? idx != res.idxs[s].back() + 1
|
||||
: idx <= res.idxs[s].back())) {
|
||||
LLAMA_LOG_WARN("%s: MSA: token positions are not %s within the ubatch\n",
|
||||
__func__, cont ? "contiguous" : "strictly increasing");
|
||||
return { };
|
||||
}
|
||||
|
||||
res.idxs[s].push_back(idx);
|
||||
}
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
uint32_t head_cur = v_heads[seq_to_stream[seq_id]];
|
||||
|
||||
// if we have enough unused cells before the current head ->
|
||||
@@ -1002,11 +1104,6 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
|
||||
head_cur = 0;
|
||||
}
|
||||
|
||||
if (n_tokens > cells.size()) {
|
||||
LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());
|
||||
return { };
|
||||
}
|
||||
|
||||
uint32_t n_tested = 0;
|
||||
|
||||
// for continuous slots, we test that all tokens in the ubatch fit, starting from the current head
|
||||
@@ -1113,6 +1210,15 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
|
||||
|
||||
const auto idx = sinfo.idxs[s][ii];
|
||||
|
||||
if (msa_strict_slots && (llama_pos) idx != ubatch.pos[i]) {
|
||||
LLAMA_LOG_ERROR("%s: MSA slot/position invariant violated: "
|
||||
"writing pos %d into cell %u (stream %u). The indexer cache "
|
||||
"would desync and block selection would silently corrupt. "
|
||||
"This is a bug, please report it with reproduction steps.\n",
|
||||
__func__, ubatch.pos[i], idx, sinfo.strm[s]);
|
||||
GGML_ABORT("MSA: slot != pos");
|
||||
}
|
||||
|
||||
if (!cells.is_empty(idx)) {
|
||||
assert(cells.seq_count(idx) == 1);
|
||||
|
||||
@@ -1156,7 +1262,8 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
|
||||
LLAMA_LOG_DEBUG("%s: purging positions [%d, %d] of sequence %d from KV cache\n",
|
||||
__func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s);
|
||||
|
||||
seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1);
|
||||
// under MSA strict slots this path should be unreachable, since strict MSA placement never selects occupied cells
|
||||
GGML_ASSERT(seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1176,6 +1283,12 @@ bool llama_kv_cache::get_can_shift() const {
|
||||
if (hparams.n_pos_per_embd() > 1) {
|
||||
return false;
|
||||
}
|
||||
// shifting would leave k_idx stale
|
||||
for (const auto & layer : layers) {
|
||||
if (layer.k_idx) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -1292,6 +1405,23 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k
|
||||
ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache::get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
|
||||
const int32_t ikv = map_layer_ids.at(il);
|
||||
auto * k_idx = layers[ikv].k_idx;
|
||||
GGML_ASSERT(k_idx);
|
||||
|
||||
const uint64_t kv_size = get_size();
|
||||
const int64_t n_idx = k_idx->ne[0]; // 128
|
||||
const uint32_t ns = sinfo.s1 - sinfo.s0 + 1;
|
||||
|
||||
return ggml_view_4d(ctx, k_idx,
|
||||
n_idx, 1, n_kv, ns,
|
||||
ggml_row_size(k_idx->type, n_idx), // nb1 (single head)
|
||||
ggml_row_size(k_idx->type, n_idx), // nb2 (per cell)
|
||||
ggml_row_size(k_idx->type, n_idx*kv_size), // nb3 (per stream)
|
||||
ggml_row_size(k_idx->type, n_idx*kv_size)*sinfo.s0);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
|
||||
GGML_UNUSED(sinfo);
|
||||
|
||||
@@ -1393,6 +1523,28 @@ ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama
|
||||
return k_idxs;
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
|
||||
GGML_UNUSED(sinfo);
|
||||
const int32_t ikv = map_layer_ids.at(il);
|
||||
ggml_tensor * k_idx = layers[ikv].k_idx;
|
||||
GGML_ASSERT(k_idx && "cpy_k_idx on a layer with no indexer cache");
|
||||
|
||||
const int64_t n_embd_head = k_idx_cur->ne[0]; // 128
|
||||
const int64_t n_head = k_idx_cur->ne[1]; // 1
|
||||
const int64_t n_tokens = k_idx_cur->ne[2];
|
||||
const int64_t n_embd_gqa = n_embd_head*n_head; // 128
|
||||
|
||||
GGML_ASSERT(ggml_row_size(k_idx_cur->type, n_embd_head) == k_idx_cur->nb[1]);
|
||||
k_idx_cur = ggml_view_2d(ctx, k_idx_cur, n_embd_gqa, n_tokens, k_idx_cur->nb[2], 0);
|
||||
|
||||
const int64_t n_stream = k_idx->ne[2];
|
||||
if (n_stream > 1) {
|
||||
const int64_t kv_size = get_size();
|
||||
k_idx = ggml_reshape_2d(ctx, k_idx, n_embd_gqa, kv_size*n_stream);
|
||||
}
|
||||
return ggml_set_rows(ctx, k_idx, k_idx_cur, k_idxs); // same k_idxs as the K store
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
|
||||
const uint32_t n_tokens = ubatch.n_tokens;
|
||||
|
||||
@@ -1827,6 +1979,18 @@ size_t llama_kv_cache::size_v_bytes() const {
|
||||
return size_v_bytes;
|
||||
}
|
||||
|
||||
size_t llama_kv_cache::size_k_idx_bytes() const {
|
||||
size_t size_k_idx_bytes = 0;
|
||||
|
||||
for (const auto & layer : layers) {
|
||||
if (layer.k_idx) {
|
||||
size_k_idx_bytes += ggml_nbytes(layer.k_idx);
|
||||
}
|
||||
}
|
||||
|
||||
return size_k_idx_bytes;
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache::build_rope_shift(
|
||||
const llama_cparams & cparams,
|
||||
ggml_context * ctx,
|
||||
@@ -2054,7 +2218,12 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama
|
||||
|
||||
bool res = true;
|
||||
res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id);
|
||||
res = res && state_read_data(io, strm, cell_count, sinfo);
|
||||
|
||||
try {
|
||||
res = res && state_read_data(io, strm, cell_count, sinfo);
|
||||
} catch (...) {
|
||||
res = false;
|
||||
}
|
||||
|
||||
if (!res) {
|
||||
if (seq_id == -1) {
|
||||
@@ -2134,6 +2303,36 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t
|
||||
}
|
||||
}
|
||||
|
||||
if (size_k_idx_bytes() > 0) {
|
||||
const uint32_t has_k_idx_u32 = 1;
|
||||
io.write(&has_k_idx_u32, sizeof(has_k_idx_u32));
|
||||
|
||||
for (const auto & layer : layers) {
|
||||
const uint32_t layer_has_k_idx = layer.k_idx ? 1 : 0;
|
||||
io.write(&layer_has_k_idx, sizeof(layer_has_k_idx));
|
||||
|
||||
if (!layer_has_k_idx) {
|
||||
continue;
|
||||
}
|
||||
|
||||
GGML_ASSERT(layer.k_idx_stream[cr.strm]);
|
||||
|
||||
const int32_t k_idx_type_i = (int32_t) layer.k_idx->type;
|
||||
io.write(&k_idx_type_i, sizeof(k_idx_type_i));
|
||||
|
||||
const uint64_t k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]);
|
||||
io.write(&k_idx_size_row, sizeof(k_idx_size_row));
|
||||
|
||||
for (const auto & range : cr.data) {
|
||||
const size_t range_size = range.second - range.first;
|
||||
const size_t buf_size = range_size * k_idx_size_row;
|
||||
const size_t offset = range.first * k_idx_size_row;
|
||||
|
||||
io.write_tensor(layer.k_idx_stream[cr.strm], offset, buf_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!v_trans) {
|
||||
for (const auto & layer : layers) {
|
||||
const uint32_t il = layer.il;
|
||||
@@ -2382,6 +2581,68 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
|
||||
}
|
||||
}
|
||||
|
||||
if (size_k_idx_bytes() > 0) {
|
||||
uint32_t has_k_idx_u32 = 0;
|
||||
io.read(&has_k_idx_u32, sizeof(has_k_idx_u32));
|
||||
|
||||
if (has_k_idx_u32 != 1) {
|
||||
LLAMA_LOG_ERROR("%s: missing k_idx data in KV cache state\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const auto & layer : layers) {
|
||||
uint32_t layer_has_k_idx = 0;
|
||||
io.read(&layer_has_k_idx, sizeof(layer_has_k_idx));
|
||||
|
||||
const uint32_t expected_layer_has_k_idx = layer.k_idx ? 1 : 0;
|
||||
|
||||
if (layer_has_k_idx != expected_layer_has_k_idx) {
|
||||
LLAMA_LOG_ERROR(
|
||||
"%s: mismatched k_idx state for layer: got %u, expected %u\n",
|
||||
__func__, layer_has_k_idx, expected_layer_has_k_idx);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!layer_has_k_idx) {
|
||||
continue;
|
||||
}
|
||||
|
||||
GGML_ASSERT(layer.k_idx_stream[strm]);
|
||||
|
||||
int32_t k_idx_type_i = -1;
|
||||
io.read(&k_idx_type_i, sizeof(k_idx_type_i));
|
||||
|
||||
if (k_idx_type_i != (int32_t) layer.k_idx->type) {
|
||||
LLAMA_LOG_ERROR(
|
||||
"%s: mismatched k_idx type: got %d, expected %d\n",
|
||||
__func__, k_idx_type_i, (int32_t) layer.k_idx->type);
|
||||
return false;
|
||||
}
|
||||
|
||||
uint64_t k_idx_size_row = 0;
|
||||
io.read(&k_idx_size_row, sizeof(k_idx_size_row));
|
||||
|
||||
const uint64_t expected_k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]);
|
||||
|
||||
if (k_idx_size_row != expected_k_idx_size_row) {
|
||||
LLAMA_LOG_ERROR(
|
||||
"%s: mismatched k_idx row size: got %zu, expected %zu\n",
|
||||
__func__, (size_t) k_idx_size_row, (size_t) expected_k_idx_size_row);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (cell_count) {
|
||||
if (sinfo.is_contiguous()) {
|
||||
io.read_tensor(layer.k_idx_stream[strm], sinfo.head() * k_idx_size_row, cell_count * k_idx_size_row);
|
||||
} else {
|
||||
for (uint32_t i = 0; i < cell_count; ++i) {
|
||||
io.read_tensor(layer.k_idx_stream[strm], sinfo.idxs[0][i] * k_idx_size_row, k_idx_size_row);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!this->v_trans) {
|
||||
for (const auto & layer : layers) {
|
||||
const uint32_t il = layer.il;
|
||||
@@ -2583,6 +2844,10 @@ ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) cons
|
||||
return kv->get_v(ctx, il, n_kv, sinfos[i_cur]);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache_context::get_k_idx(ggml_context * ctx, int32_t il) const {
|
||||
return kv->get_k_idx(ctx, il, n_kv, sinfos[i_cur]);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const {
|
||||
return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]);
|
||||
}
|
||||
@@ -2591,6 +2856,10 @@ ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_
|
||||
return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache_context::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const {
|
||||
return kv->cpy_k_idx(ctx, k_idx_cur, k_idxs, il, sinfos[i_cur]);
|
||||
}
|
||||
|
||||
ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const {
|
||||
return kv->build_input_k_idxs(ctx, ubatch);
|
||||
}
|
||||
|
||||
@@ -173,10 +173,12 @@ public:
|
||||
// get views of the current state of the cache
|
||||
ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
|
||||
ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
|
||||
ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
|
||||
|
||||
// store k_cur and v_cur in the cache based on the provided head location
|
||||
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
|
||||
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const;
|
||||
ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
|
||||
|
||||
//
|
||||
// preparation API
|
||||
@@ -228,9 +230,11 @@ private:
|
||||
|
||||
ggml_tensor * k;
|
||||
ggml_tensor * v;
|
||||
ggml_tensor * k_idx; // MSA single-head indexer keys, F32
|
||||
|
||||
std::vector<ggml_tensor *> k_stream;
|
||||
std::vector<ggml_tensor *> v_stream;
|
||||
std::vector<ggml_tensor *> k_idx_stream;
|
||||
};
|
||||
|
||||
bool v_trans = true; // the value tensor is transposed
|
||||
@@ -259,6 +263,9 @@ private:
|
||||
// env: LLAMA_KV_CACHE_DEBUG
|
||||
int debug = 0;
|
||||
|
||||
// set when a k_idx (indexer) cache exists and the stream layout supports MSA (single seq, or one stream per seq)
|
||||
bool msa_strict_slots = false;
|
||||
|
||||
// this is the SWA type of the cache - not to be confused with the model SWA type
|
||||
const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
|
||||
|
||||
@@ -291,6 +298,7 @@ private:
|
||||
|
||||
size_t size_k_bytes() const;
|
||||
size_t size_v_bytes() const;
|
||||
size_t size_k_idx_bytes() const;
|
||||
|
||||
ggml_tensor * build_rope_shift(
|
||||
const llama_cparams & cparams,
|
||||
@@ -370,6 +378,7 @@ public:
|
||||
// get views of the current state of the cache
|
||||
ggml_tensor * get_k(ggml_context * ctx, int32_t il) const;
|
||||
ggml_tensor * get_v(ggml_context * ctx, int32_t il) const;
|
||||
ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il) const;
|
||||
|
||||
// store k_cur and v_cur in the cache based on the provided head location
|
||||
// note: the heads in k_cur and v_cur should be laid out contiguously in memory
|
||||
@@ -379,6 +388,7 @@ public:
|
||||
// - v_idxs [n_tokens] or [n_tokens*n_embd_v_gqa] depending if V cache is transposed
|
||||
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const;
|
||||
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const;
|
||||
ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const;
|
||||
|
||||
// create destination indices for each head of the current batch for where it would be written in the KV cache
|
||||
// the indices address the global KV cache (not per stream) - this is not relevant for the user of this API, but
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user