mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-08-04 17:08:38 +00:00
Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2f56fc3431 | ||
|
|
0713275082 | ||
|
|
1c3c9674de | ||
|
|
6b5224cfcc |
@@ -63,7 +63,8 @@ jobs:
|
||||
-DGGML_METAL_USE_BF16=ON \
|
||||
-DGGML_METAL_EMBED_LIBRARY=OFF \
|
||||
-DGGML_METAL_SHADER_DEBUG=ON \
|
||||
-DGGML_RPC=ON
|
||||
-DGGML_RPC=ON \
|
||||
-DCMAKE_OSX_DEPLOYMENT_TARGET=13.3
|
||||
time cmake --build build --config Release -j $(sysctl -n hw.logicalcpu)
|
||||
leaks -atExit -- ./build/bin/test-thread-safety -hf ggml-org/gemma-3-270m-qat-GGUF -ngl 99 -p "$(printf 'hello %.0s' {1..128})" -n 16 -c 512 -ub 32 -np 2 -t 2 -lv 1
|
||||
|
||||
|
||||
@@ -93,13 +93,13 @@ jobs:
|
||||
- build: 'arm64'
|
||||
arch: 'arm64'
|
||||
os: macos-26
|
||||
defines: "-DGGML_METAL_USE_BF16=ON -DGGML_METAL_EMBED_LIBRARY=ON"
|
||||
defines: "-DGGML_METAL_USE_BF16=ON -DGGML_METAL_EMBED_LIBRARY=ON -DCMAKE_OSX_DEPLOYMENT_TARGET=13.3"
|
||||
# TODO: this build is disabled to save Github Actions resources (https://github.com/ggml-org/llama.cpp/pull/23780)
|
||||
# in order to enable it again, we have to provision dedicated runners to run it
|
||||
#- build: 'arm64-kleidiai'
|
||||
# arch: 'arm64'
|
||||
# os: macos-14
|
||||
# defines: "-DGGML_METAL_USE_BF16=ON -DGGML_METAL_EMBED_LIBRARY=ON -DGGML_CPU_KLEIDIAI=ON"
|
||||
# defines: "-DGGML_METAL_USE_BF16=ON -DGGML_METAL_EMBED_LIBRARY=ON -DCMAKE_OSX_DEPLOYMENT_TARGET=13.3 -DGGML_CPU_KLEIDIAI=ON"
|
||||
- build: 'x64'
|
||||
arch: 'x64'
|
||||
os: macos-15-intel
|
||||
|
||||
+12
-57
@@ -61,6 +61,7 @@ static std::initializer_list<enum llama_example> mmproj_examples = {
|
||||
LLAMA_EXAMPLE_MTMD,
|
||||
LLAMA_EXAMPLE_SERVER,
|
||||
LLAMA_EXAMPLE_CLI,
|
||||
LLAMA_EXAMPLE_TTS,
|
||||
};
|
||||
|
||||
static std::string read_file(const std::string & fname) {
|
||||
@@ -360,7 +361,6 @@ static bool spec_types_is_default(const common_params & params) {
|
||||
common_models_handler common_models_handler_init(const common_params & params, llama_example curr_ex) {
|
||||
common_download_hf_plan plan;
|
||||
common_download_hf_plan plan_spec;
|
||||
common_download_hf_plan plan_voc;
|
||||
common_download_opts opts;
|
||||
|
||||
const bool spec_type_draft_mtp = std::find(params.speculative.types.begin(),
|
||||
@@ -413,11 +413,7 @@ common_models_handler common_models_handler_init(const common_params & params, l
|
||||
plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts_spec);
|
||||
}
|
||||
|
||||
if (!params.vocoder.model.hf_repo.empty()) {
|
||||
plan_voc = common_download_get_hf_plan(params.vocoder.model, opts);
|
||||
}
|
||||
|
||||
return common_models_handler{plan, plan_spec, plan_voc, opts};
|
||||
return common_models_handler{plan, plan_spec, opts};
|
||||
}
|
||||
|
||||
bool common_models_handler_is_preset_repo(const common_models_handler & handler) {
|
||||
@@ -467,7 +463,6 @@ void common_models_handler_apply(common_models_handler & handler, common_params
|
||||
|
||||
auto & plan = handler.plan;
|
||||
auto & plan_spec = handler.plan_spec;
|
||||
auto & plan_voc = handler.plan_voc;
|
||||
|
||||
auto opts = handler.opts; // copy
|
||||
opts.callback = callback;
|
||||
@@ -482,7 +477,6 @@ void common_models_handler_apply(common_models_handler & handler, common_params
|
||||
};
|
||||
handle_url(params.model);
|
||||
handle_url(params.mmproj);
|
||||
handle_url(params.vocoder.model);
|
||||
handle_url(params.speculative.draft.mparams);
|
||||
|
||||
// optionally, if docker repo is set, resolve it
|
||||
@@ -510,14 +504,6 @@ void common_models_handler_apply(common_models_handler & handler, common_params
|
||||
task.opts = opts;
|
||||
tasks.push_back(task);
|
||||
}
|
||||
if (!params.vocoder.model.url.empty()) {
|
||||
common_download_task task;
|
||||
task.url = params.vocoder.model.url;
|
||||
task.local_path = params.vocoder.model.path;
|
||||
task.opts = opts;
|
||||
tasks.push_back(task);
|
||||
}
|
||||
|
||||
bool had_spec_url = false;
|
||||
if (!params.speculative.draft.mparams.url.empty()) {
|
||||
common_download_task task;
|
||||
@@ -631,11 +617,6 @@ void common_models_handler_apply(common_models_handler & handler, common_params
|
||||
had_spec_url = true;
|
||||
}
|
||||
|
||||
// handle vocoder plan (e.g. --hf-repo-v)
|
||||
if (!plan_voc.model_files.empty()) {
|
||||
add_tasks(plan_voc.model_files, plan_voc.primary, params.vocoder.model);
|
||||
}
|
||||
|
||||
if (!plan.model_files.empty()) {
|
||||
add_tasks(plan.model_files, plan.primary, params.model);
|
||||
}
|
||||
@@ -1361,6 +1342,10 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
params.n_parallel = -1; // auto by default
|
||||
} else if (ex == LLAMA_EXAMPLE_TOKENIZE) {
|
||||
params.parse_special = true; // parse special tokens by default, like the old tokenize tool
|
||||
} else if (ex == LLAMA_EXAMPLE_TTS) {
|
||||
params.out_file = "output.wav";
|
||||
params.sampling.penalty_repeat = 1.05f;
|
||||
params.sampling.penalty_last_n = -1;
|
||||
}
|
||||
|
||||
params.use_color = tty_can_use_colors();
|
||||
@@ -2983,20 +2968,6 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
params.model.hf_file = value;
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_TOKENIZE}).set_env("LLAMA_ARG_HF_FILE"));
|
||||
add_opt(common_arg(
|
||||
{"-hfv", "-hfrv", "--hf-repo-v"}, "<user>/<model>[:quant]",
|
||||
"Hugging Face model repository for the vocoder model (default: unused)",
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.vocoder.model.hf_repo = value;
|
||||
}
|
||||
).set_env("LLAMA_ARG_HF_REPO_V"));
|
||||
add_opt(common_arg(
|
||||
{"-hffv", "--hf-file-v"}, "FILE",
|
||||
"Hugging Face model file for the vocoder model (default: unused)",
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.vocoder.model.hf_file = value;
|
||||
}
|
||||
).set_env("LLAMA_ARG_HF_FILE_V"));
|
||||
add_opt(common_arg(
|
||||
{"-hft", "--hf-token"}, "TOKEN",
|
||||
"Hugging Face access token (default: value from HF_TOKEN environment variable)",
|
||||
@@ -4272,24 +4243,18 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
//
|
||||
|
||||
add_opt(common_arg(
|
||||
{"-mv", "--model-vocoder"}, "FNAME",
|
||||
"vocoder model for audio generation (default: unused)",
|
||||
{"--tts-lang"}, "FNAME",
|
||||
"language (ISO 639-1) for audio generation\n"
|
||||
"see tts/README.md for per-model usage notes",
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.vocoder.model.path = value;
|
||||
params.tts_lang = value;
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_SERVER}));
|
||||
add_opt(common_arg(
|
||||
{"--tts-use-guide-tokens"},
|
||||
"Use guide tokens to improve TTS word recall",
|
||||
[](common_params & params) {
|
||||
params.vocoder.use_guide_tokens = true;
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_SERVER}));
|
||||
).set_examples({LLAMA_EXAMPLE_TTS}));
|
||||
add_opt(common_arg(
|
||||
{"--tts-speaker-file"}, "FNAME",
|
||||
"speaker file path for audio generation",
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.vocoder.speaker_file = value;
|
||||
params.tts_speaker_file = value;
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_TTS}));
|
||||
|
||||
@@ -4409,16 +4374,6 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
).set_examples({LLAMA_EXAMPLE_DEBUG}));
|
||||
|
||||
// presets
|
||||
add_opt(common_arg(
|
||||
{"--tts-oute-default"},
|
||||
string_format("use default OuteTTS models (note: can download weights from the internet)"),
|
||||
[](common_params & params) {
|
||||
params.model.hf_repo = "OuteAI/OuteTTS-0.2-500M-GGUF";
|
||||
params.model.hf_file = "OuteTTS-0.2-500M-Q8_0.gguf";
|
||||
params.vocoder.model.hf_repo = "ggml-org/WavTokenizer";
|
||||
params.vocoder.model.hf_file = "WavTokenizer-Large-75-F16.gguf";
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_TTS}));
|
||||
|
||||
add_opt(common_arg(
|
||||
{"--embd-gemma-default"},
|
||||
|
||||
@@ -137,7 +137,6 @@ void common_params_add_preset_options(std::vector<common_arg> & args);
|
||||
struct common_models_handler {
|
||||
common_download_hf_plan plan;
|
||||
common_download_hf_plan plan_spec;
|
||||
common_download_hf_plan plan_voc;
|
||||
common_download_opts opts;
|
||||
};
|
||||
|
||||
|
||||
+4
-9
@@ -392,14 +392,6 @@ struct common_params_speculative {
|
||||
}
|
||||
};
|
||||
|
||||
struct common_params_vocoder {
|
||||
struct common_params_model model;
|
||||
|
||||
std::string speaker_file; // speaker file path
|
||||
|
||||
bool use_guide_tokens = false; // enable guide tokens to improve TTS accuracy
|
||||
};
|
||||
|
||||
struct common_params_diffusion {
|
||||
int32_t steps = 128;
|
||||
bool visual_mode = false;
|
||||
@@ -497,7 +489,6 @@ struct common_params {
|
||||
|
||||
struct common_params_sampling sampling;
|
||||
struct common_params_speculative speculative;
|
||||
struct common_params_vocoder vocoder;
|
||||
struct common_params_diffusion diffusion;
|
||||
|
||||
struct common_params_model model;
|
||||
@@ -740,6 +731,10 @@ struct common_params {
|
||||
void * load_progress_callback_user_data = NULL;
|
||||
bool no_alloc = false; // Don't allocate model buffers
|
||||
|
||||
// TTS params
|
||||
std::string tts_lang = "";
|
||||
std::string tts_speaker_file = "";
|
||||
|
||||
bool is_gen_docs = false; // whether we are running inside llama-gen-docs
|
||||
};
|
||||
|
||||
|
||||
@@ -210,6 +210,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"Qwen3MoeForCausalLM": "qwen",
|
||||
"Qwen3NextForCausalLM": "qwen",
|
||||
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
|
||||
"Qwen3TTSForConditionalGeneration": "qwen3tts",
|
||||
"Qwen3VLForConditionalGeneration": "qwen3vl",
|
||||
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
|
||||
"Qwen3_5ForCausalLM": "qwen",
|
||||
@@ -304,6 +305,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
|
||||
"Qwen2_5_VLForConditionalGeneration": "qwenvl",
|
||||
"Qwen3ASRForConditionalGeneration": "qwen3vl",
|
||||
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
|
||||
"Qwen3TTSForConditionalGeneration": "qwen3tts",
|
||||
"Qwen3VLForConditionalGeneration": "qwen3vl",
|
||||
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
|
||||
"Qwen3_5ForConditionalGeneration": "qwen3vl",
|
||||
|
||||
@@ -0,0 +1,471 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Iterable, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch import Tensor
|
||||
|
||||
from .base import ModelBase, MmprojModel, TextModel, gguf
|
||||
|
||||
# Tricks being used to support this model via existing llama.cpp code paths:
|
||||
# - Text projection MLP is folded into the embedding table
|
||||
# - codec_embedding is concat to the text embedding table, vocab is extended
|
||||
# example: codec_bos_id(2149) --> "<|codec_bos|>"
|
||||
# codec_eos_token_id(2150) --> "<|codec_eos_token|>"
|
||||
# codec_language_id.chinese(2055) --> "<|codec_language_chinese|>"
|
||||
# other rows --> "<|codec_0|>", "<|codec_1|>", ..., "<|codec_1023|>"
|
||||
# - output tensor codec_head is smaller than vocab, so logits will be padded at inference time
|
||||
# - suppress_tokens is used to limit the backbone to only sample either semantic or EOS (stop) token
|
||||
|
||||
# pipeline stage mapping:
|
||||
# speaker reference encoder --> mapped to normal mtmd audio encoder
|
||||
# backbone --> mapped to normal libllama text model (autoregressive)
|
||||
# code_predictor --> MTMD_GEN_PROCESS_TYPE_GEN_CODE
|
||||
# code2wav --> MTMD_GEN_PROCESS_TYPE_GEN_WAV
|
||||
|
||||
# torch activation functions used by Qwen3TTSTalkerResizeMLP (config's hidden_act)
|
||||
_ACT2FN = {
|
||||
"silu": F.silu,
|
||||
"gelu": F.gelu,
|
||||
"relu": F.relu,
|
||||
}
|
||||
|
||||
|
||||
@ModelBase.register("Qwen3TTSForConditionalGeneration")
|
||||
class Qwen3TTSTalkerModel(TextModel):
|
||||
model_arch = gguf.MODEL_ARCH.QWEN3TTS
|
||||
|
||||
_TEXT_PROJ_KEYS = (
|
||||
"model.text_embedding.weight",
|
||||
"text_projection.linear_fc1.weight",
|
||||
"text_projection.linear_fc1.bias",
|
||||
"text_projection.linear_fc2.weight",
|
||||
"text_projection.linear_fc2.bias",
|
||||
)
|
||||
|
||||
_text_proj_buffer: dict[str, Tensor]
|
||||
_folded_text_embed: Tensor | None
|
||||
_codec_embed: Tensor | None
|
||||
|
||||
def __init__(self, dir_model: Path, *args, **kwargs):
|
||||
hparams = kwargs.pop("hparams", None)
|
||||
if hparams is None:
|
||||
hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
|
||||
raw_talker_config = dict(hparams["talker_config"])
|
||||
self._talker_config = raw_talker_config
|
||||
self.n_codec_vocab = raw_talker_config["vocab_size"]
|
||||
talker_config = dict(raw_talker_config)
|
||||
talker_config["vocab_size"] = talker_config["text_vocab_size"]
|
||||
hparams["text_config"] = talker_config
|
||||
super().__init__(dir_model, *args, hparams=hparams, **kwargs)
|
||||
self._text_proj_buffer = {}
|
||||
self._folded_text_embed = None
|
||||
self._codec_embed = None
|
||||
|
||||
def _codec_token_names(self) -> list[str]:
|
||||
# start every row with a generic name, then override the ones with a
|
||||
# known meaning (bos/eos/language/etc, derived from the *_id fields
|
||||
# of talker_config) with a more descriptive one
|
||||
names = [f"<|codec_{i}|>" for i in range(self.n_codec_vocab)]
|
||||
for key, val in self._talker_config.items():
|
||||
if not key.endswith("_id"):
|
||||
continue
|
||||
prefix = key[:-len("_id")]
|
||||
if isinstance(val, int):
|
||||
names[val] = f"<|{prefix}|>"
|
||||
elif isinstance(val, dict):
|
||||
for subkey, subval in val.items():
|
||||
names[subval] = f"<|{prefix}_{subkey}|>"
|
||||
return names
|
||||
|
||||
def set_vocab(self):
|
||||
codec_tokens = self._codec_token_names()
|
||||
codec_toktypes = [gguf.TokenType.CONTROL] * len(codec_tokens)
|
||||
|
||||
try:
|
||||
tokens, scores, toktypes = self._create_vocab_sentencepiece()
|
||||
self.gguf_writer.add_tokenizer_model("llama")
|
||||
self.gguf_writer.add_tokenizer_pre("default")
|
||||
tokens += [t.encode("utf-8") for t in codec_tokens]
|
||||
scores += [0.0] * len(codec_tokens)
|
||||
toktypes += codec_toktypes
|
||||
self.gguf_writer.add_token_list(tokens)
|
||||
self.gguf_writer.add_token_scores(scores)
|
||||
self.gguf_writer.add_token_types(toktypes)
|
||||
special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))
|
||||
special_vocab.add_to_gguf(self.gguf_writer)
|
||||
return
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
tokens, toktypes, tokpre = self.get_vocab_base()
|
||||
tokens += codec_tokens
|
||||
toktypes += codec_toktypes
|
||||
self.gguf_writer.add_tokenizer_model("gpt2")
|
||||
self.gguf_writer.add_tokenizer_pre(tokpre)
|
||||
self.gguf_writer.add_token_list(tokens)
|
||||
self.gguf_writer.add_token_types(toktypes)
|
||||
|
||||
special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
|
||||
special_vocab.add_to_gguf(self.gguf_writer)
|
||||
|
||||
# make sure that the model has no chat template, so chat will be disabled
|
||||
self.gguf_writer.add_chat_template(None)
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
super().set_gguf_parameters()
|
||||
|
||||
# note: final vocab layout is [text_vocab | codec_vocab], with text_vocab is actually padded with -inf in cgraph
|
||||
# for codec_vocab, only first 2048 rows can be sampled for semantic code
|
||||
# plus codec_eos_token_id that used for signaling end of generation
|
||||
# ref: https://github.com/QwenLM/Qwen3-TTS/blob/022e286b98fbec7e1e916cb940cdf532cd9f488e/qwen_tts/core/models/modeling_qwen3_tts.py#L2059-L2063
|
||||
|
||||
vocab_size = self.hparams["vocab_size"] + self.n_codec_vocab
|
||||
codec_eos_token_id = self.hparams["vocab_size"] + self._talker_config["codec_eos_token_id"]
|
||||
self.gguf_writer.add_suppress_tokens([
|
||||
i for i in range(vocab_size - 1024, vocab_size)
|
||||
if i != codec_eos_token_id
|
||||
])
|
||||
self.gguf_writer.add_eos_token_id(codec_eos_token_id)
|
||||
self.gguf_writer.add_add_eos_token(False)
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
name, gen = item
|
||||
|
||||
if not name.startswith("talker.") or name.startswith("talker.code_predictor."):
|
||||
return None
|
||||
|
||||
name = name[len("talker."):]
|
||||
return super().filter_tensors((name, gen))
|
||||
|
||||
def _maybe_emit_token_embd(self) -> Iterable[tuple[str, Tensor]]:
|
||||
if self._folded_text_embed is None or self._codec_embed is None:
|
||||
return
|
||||
combined = torch.cat([self._folded_text_embed, self._codec_embed], dim=0)
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), combined)
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
# codec_embedding rows are appended after the text vocab, extending the embedding table
|
||||
if name == "model.codec_embedding.weight":
|
||||
self._codec_embed = data_torch
|
||||
yield from self._maybe_emit_token_embd()
|
||||
return
|
||||
|
||||
# codec_head is the output head for the (smaller) codec vocab; logits get padded to
|
||||
# the extended vocab size at inference time
|
||||
if name == "codec_head.weight":
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT), data_torch)
|
||||
return
|
||||
|
||||
if name in self._TEXT_PROJ_KEYS:
|
||||
self._text_proj_buffer[name] = data_torch
|
||||
if len(self._text_proj_buffer) < len(self._TEXT_PROJ_KEYS):
|
||||
return
|
||||
|
||||
# fold MLP into the embedding table at conversion time, MLP won't be used at inference time anyway
|
||||
act_fn = _ACT2FN[self.hparams["hidden_act"]]
|
||||
embed = self._text_proj_buffer["model.text_embedding.weight"]
|
||||
hidden = act_fn(F.linear(embed,
|
||||
self._text_proj_buffer["text_projection.linear_fc1.weight"],
|
||||
self._text_proj_buffer["text_projection.linear_fc1.bias"]))
|
||||
folded = F.linear(hidden,
|
||||
self._text_proj_buffer["text_projection.linear_fc2.weight"],
|
||||
self._text_proj_buffer["text_projection.linear_fc2.bias"])
|
||||
self._folded_text_embed = folded
|
||||
yield from self._maybe_emit_token_embd()
|
||||
return
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
|
||||
@ModelBase.register("Qwen3TTSForConditionalGeneration")
|
||||
class Qwen3TTSSpeakerEncoderModel(MmprojModel):
|
||||
has_vision_encoder = False
|
||||
has_audio_encoder = True
|
||||
|
||||
# talker.code_predictor.model.layers.{bid}.<key> -> A_GEN_CODE_*
|
||||
# bypass tensor_mapping.py for now to make it simple
|
||||
_CODE_LAYER_TENSOR_MAP = {
|
||||
"input_layernorm": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_NORM,
|
||||
"self_attn.q_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_Q,
|
||||
"self_attn.q_norm": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM,
|
||||
"self_attn.k_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_K,
|
||||
"self_attn.k_norm": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM,
|
||||
"self_attn.v_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_V,
|
||||
"self_attn.o_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_OUT,
|
||||
"post_attention_layernorm": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_NORM,
|
||||
"mlp.gate_proj": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_GATE,
|
||||
"mlp.up_proj": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_UP,
|
||||
"mlp.down_proj": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_DOWN,
|
||||
}
|
||||
|
||||
# note: codebook pages will be stacked to 3D
|
||||
_CODE_GEN_N_CODEBOOKS = 15
|
||||
_code_embed_buffer: dict[int, Tensor] = {}
|
||||
_code_head_buffer: dict[int, Tensor] = {}
|
||||
_wav_config_cache: dict[str, Any] | None = None
|
||||
|
||||
def __init__(self, dir_model: Path, *args, **kwargs):
|
||||
hparams = kwargs.pop("hparams", None)
|
||||
if hparams is None:
|
||||
hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
|
||||
hparams["text_config"] = {"hidden_size": hparams["talker_config"]["hidden_size"]}
|
||||
# ECAPA-TDNN has a fixed 4-stage backbone, but MmprojModel.__init__ needs a n_block_keys
|
||||
hparams["speaker_encoder_config"]["n_layers"] = 4
|
||||
super().__init__(dir_model, *args, hparams=hparams, **kwargs)
|
||||
self._wav_config_cache = None
|
||||
|
||||
def get_audio_config(self) -> dict[str, Any] | None:
|
||||
return self.global_config.get("speaker_encoder_config")
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
self.gguf_writer.add_file_type(self.ftype)
|
||||
self.gguf_writer.add_clip_has_audio_encoder(True)
|
||||
self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.QWEN3TTS_SPKENC)
|
||||
|
||||
# handle speaker encoder config
|
||||
self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
|
||||
# mel_spectrogram() front-end: sr=24000, n_fft=1024, hop=256, n_mels=128, fmin=0, fmax=12000 (=sr/2, the clip.cpp default)
|
||||
self.gguf_writer.add_audio_num_mel_bins(128)
|
||||
# 3 SE-Res2Net stages; the stem conv, mfa, asp and fc are not counted here
|
||||
self.gguf_writer.add_audio_block_count(3)
|
||||
# ECAPA-TDNN has no attention/FFN, these are dummy to allow clip.cpp to load it
|
||||
self.gguf_writer.add_audio_embedding_length(1536)
|
||||
self.gguf_writer.add_audio_head_count(1)
|
||||
self.gguf_writer.add_audio_feed_forward_length(1536)
|
||||
self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
|
||||
|
||||
# handle code predictor config
|
||||
self.gguf_writer.add_clip_has_gen_audio_encoder(True)
|
||||
self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.QWEN3TTS_GEN)
|
||||
code_predictor_config = self.global_config["talker_config"]["code_predictor_config"]
|
||||
self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)
|
||||
self.gguf_writer.add_gen_audio_embedding_length(code_predictor_config["hidden_size"])
|
||||
self.gguf_writer.add_gen_audio_feed_forward_length(code_predictor_config["intermediate_size"])
|
||||
self.gguf_writer.add_gen_audio_block_count(code_predictor_config["num_hidden_layers"])
|
||||
self.gguf_writer.add_gen_audio_head_count(code_predictor_config["num_attention_heads"])
|
||||
self.gguf_writer.add_gen_audio_head_count_kv(code_predictor_config["num_key_value_heads"])
|
||||
self.gguf_writer.add_gen_audio_attention_layernorm_eps(code_predictor_config["rms_norm_eps"])
|
||||
# note: code2wav hparams are hardcoded on the mtmd/clip.cpp side for now, not written here
|
||||
|
||||
def _wav_decoder_config(self) -> dict[str, Any] | None:
|
||||
# code2wav has its own config.json, inside the speech_tokenizer dir
|
||||
if self._wav_config_cache is None:
|
||||
path = self.dir_model / "speech_tokenizer" / "config.json"
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
self._wav_config_cache = cfg["decoder_config"]
|
||||
return self._wav_config_cache
|
||||
|
||||
def tensor_force_quant(self, name, new_name, bid, n_dims):
|
||||
# conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path
|
||||
if new_name.endswith(".weight") and (
|
||||
new_name in ("a.gen.wav.pre_conv.weight", "a.gen.wav.dac.entry.weight", "a.gen.wav.dac.post_conv.weight")
|
||||
or (".up.blk." in new_name and new_name.endswith(".dwconv.weight"))
|
||||
or (".dac.blk." in new_name and (new_name.endswith(".conv1.weight") or new_name.endswith(".conv2.weight")))
|
||||
):
|
||||
return gguf.GGMLQuantizationType.F16
|
||||
# ConvTranspose1d kernels: only F16/F32 are implemented, no BF16
|
||||
if new_name.endswith(".conv.weight") and (".up.blk." in new_name or ".dac.blk." 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 not (
|
||||
name.startswith("speaker_encoder.")
|
||||
or name.startswith("talker.code_predictor.")
|
||||
or name == "talker.model.codec_embedding.weight"
|
||||
):
|
||||
return None
|
||||
|
||||
return super().filter_tensors((name, gen))
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
# code2wav tensors are already named by generate_extra_tensors(), pass them through
|
||||
if name.startswith("a.gen.wav."):
|
||||
yield (name, data_torch)
|
||||
return
|
||||
|
||||
# codebook-0 embedding, fed back to the talker backbone (codebooks 1-15 live in code_predictor)
|
||||
if name == "talker.model.codec_embedding.weight":
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_OUT_EMBD), data_torch)
|
||||
return
|
||||
|
||||
if name == "talker.code_predictor.model.norm.weight":
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_OUTPUT_NORM), data_torch)
|
||||
return
|
||||
|
||||
if name.startswith("talker.code_predictor.small_to_mtp_projection."):
|
||||
suffix = "." + name.rsplit(".", 1)[1]
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_PROJ_IN, suffix=suffix), data_torch)
|
||||
return
|
||||
|
||||
if name.startswith("talker.code_predictor.model.codec_embedding."):
|
||||
idx = int(name.split("codec_embedding.")[1].split(".")[0])
|
||||
self._code_embed_buffer[idx] = data_torch
|
||||
if len(self._code_embed_buffer) < self._CODE_GEN_N_CODEBOOKS:
|
||||
return
|
||||
stacked = torch.stack([self._code_embed_buffer.pop(i) for i in range(self._CODE_GEN_N_CODEBOOKS)], dim=0)
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_EMBD), stacked)
|
||||
return
|
||||
|
||||
if name.startswith("talker.code_predictor.lm_head."):
|
||||
idx = int(name.split("lm_head.")[1].split(".")[0])
|
||||
self._code_head_buffer[idx] = data_torch
|
||||
if len(self._code_head_buffer) < self._CODE_GEN_N_CODEBOOKS:
|
||||
return
|
||||
stacked = torch.stack([self._code_head_buffer.pop(i) for i in range(self._CODE_GEN_N_CODEBOOKS)], dim=0)
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_HEAD), stacked)
|
||||
return
|
||||
|
||||
if name.startswith("talker.code_predictor.model.layers."):
|
||||
rest = name.split("model.layers.")[1] # "{bid}.<key>.weight"
|
||||
_, key_with_suffix = rest.split(".", 1) # "<key>.weight"
|
||||
key = key_with_suffix.rsplit(".", 1)[0] # "<key>"
|
||||
tensor = self._CODE_LAYER_TENSOR_MAP.get(key)
|
||||
if tensor is not None:
|
||||
yield (self.format_tensor_name(tensor, bid), data_torch)
|
||||
return
|
||||
|
||||
if "res2net_block.blocks." in name:
|
||||
assert bid is not None # the outer stage index, picked up from the tensor name automatically
|
||||
xid = int(name.split("res2net_block.blocks.")[1].split(".")[0])
|
||||
suffix = "." + name.rsplit(".", 1)[1]
|
||||
new_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_ENC_CONV_RES2].format(bid=bid, xid=xid) + suffix
|
||||
yield (new_name, data_torch)
|
||||
return
|
||||
|
||||
yield from super().modify_tensors(data_torch, name, bid)
|
||||
|
||||
def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
|
||||
yield from self._generate_code2wav_tensors()
|
||||
|
||||
def _generate_code2wav_tensors(self) -> Iterable[tuple[str, Tensor]]:
|
||||
# code2wav weights live in speech_tokenizer/model.safetensors, not the main safetensors
|
||||
from safetensors.torch import load_file
|
||||
|
||||
wav_config = self._wav_decoder_config()
|
||||
state_dict = load_file(self.dir_model / "speech_tokenizer" / "model.safetensors")
|
||||
|
||||
def get(name: str) -> Tensor:
|
||||
return state_dict[name]
|
||||
|
||||
def snake_fold(alpha: Tensor, beta: Tensor) -> tuple[Tensor, Tensor]:
|
||||
# fold SnakeBeta's exp()/reciprocal here, so the graph is only mul/sin/sqr/mul/add
|
||||
return torch.exp(alpha), 1.0 / (torch.exp(beta) + 1e-9)
|
||||
|
||||
def rvq_codebook(prefix: str, n_layers: int) -> Tensor:
|
||||
# checkpoint has EMA accumulators, so codebook[i] = embedding_sum[i] / cluster_usage[i]
|
||||
books = []
|
||||
for i in range(n_layers):
|
||||
embedding_sum = get(f"{prefix}.vq.layers.{i}._codebook.embedding_sum")
|
||||
cluster_usage = get(f"{prefix}.vq.layers.{i}._codebook.cluster_usage")
|
||||
books.append(embedding_sum / cluster_usage.clamp_min(1e-5).unsqueeze(-1))
|
||||
return torch.stack(books, dim=0) if n_layers > 1 else books[0]
|
||||
|
||||
T = gguf.MODEL_TENSOR
|
||||
|
||||
# --- quantizer: RVQ codebook decode ---
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_IN), get("decoder.quantizer.rvq_first.input_proj.weight").squeeze(-1))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_OUT), get("decoder.quantizer.rvq_first.output_proj.weight").squeeze(-1))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_CB), rvq_codebook("decoder.quantizer.rvq_first", 1))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_REST_IN), get("decoder.quantizer.rvq_rest.input_proj.weight").squeeze(-1))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_REST_OUT), get("decoder.quantizer.rvq_rest.output_proj.weight").squeeze(-1))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_REST_CB), rvq_codebook("decoder.quantizer.rvq_rest", self._CODE_GEN_N_CODEBOOKS))
|
||||
|
||||
# --- pre_conv ---
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_PRE_CONV, suffix=".weight"), get("decoder.pre_conv.conv.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_PRE_CONV, suffix=".bias"), get("decoder.pre_conv.conv.bias"))
|
||||
|
||||
# --- pre_transformer ---
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_IN_PROJ, suffix=".weight"), get("decoder.pre_transformer.input_proj.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_IN_PROJ, suffix=".bias"), get("decoder.pre_transformer.input_proj.bias"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUT_PROJ, suffix=".weight"), get("decoder.pre_transformer.output_proj.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUT_PROJ, suffix=".bias"), get("decoder.pre_transformer.output_proj.bias"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUTPUT_NORM), get("decoder.pre_transformer.norm.weight"))
|
||||
|
||||
tfm_layer_map = {
|
||||
"input_layernorm.weight": T.A_GEN_WAV_TFM_ATTN_NORM,
|
||||
"self_attn.q_proj.weight": T.A_GEN_WAV_TFM_ATTN_Q,
|
||||
"self_attn.k_proj.weight": T.A_GEN_WAV_TFM_ATTN_K,
|
||||
"self_attn.v_proj.weight": T.A_GEN_WAV_TFM_ATTN_V,
|
||||
"self_attn.o_proj.weight": T.A_GEN_WAV_TFM_ATTN_OUT,
|
||||
"self_attn_layer_scale.scale": T.A_GEN_WAV_TFM_ATTN_SCALE,
|
||||
"post_attention_layernorm.weight": T.A_GEN_WAV_TFM_FFN_NORM,
|
||||
"mlp.gate_proj.weight": T.A_GEN_WAV_TFM_FFN_GATE,
|
||||
"mlp.up_proj.weight": T.A_GEN_WAV_TFM_FFN_UP,
|
||||
"mlp.down_proj.weight": T.A_GEN_WAV_TFM_FFN_DOWN,
|
||||
"mlp_layer_scale.scale": T.A_GEN_WAV_TFM_FFN_SCALE,
|
||||
}
|
||||
assert wav_config is not None
|
||||
for bid in range(wav_config["num_hidden_layers"]):
|
||||
for key, tensor_id in tfm_layer_map.items():
|
||||
yield (self.format_tensor_name(tensor_id, bid), get(f"decoder.pre_transformer.layers.{bid}.{key}"))
|
||||
|
||||
# --- upsample: 2x (causal ConvTranspose1d + ConvNeXt block) ---
|
||||
up_map = {
|
||||
"0.conv.weight": (T.A_GEN_WAV_UP_CONV, ".weight"),
|
||||
"0.conv.bias": (T.A_GEN_WAV_UP_CONV, ".bias"),
|
||||
"1.dwconv.conv.weight": (T.A_GEN_WAV_UP_DWCONV, ".weight"),
|
||||
"1.dwconv.conv.bias": (T.A_GEN_WAV_UP_DWCONV, ".bias"),
|
||||
"1.norm.weight": (T.A_GEN_WAV_UP_NORM, ".weight"),
|
||||
"1.norm.bias": (T.A_GEN_WAV_UP_NORM, ".bias"),
|
||||
"1.pwconv1.weight": (T.A_GEN_WAV_UP_PW1, ".weight"),
|
||||
"1.pwconv1.bias": (T.A_GEN_WAV_UP_PW1, ".bias"),
|
||||
"1.pwconv2.weight": (T.A_GEN_WAV_UP_PW2, ".weight"),
|
||||
"1.pwconv2.bias": (T.A_GEN_WAV_UP_PW2, ".bias"),
|
||||
"1.gamma": (T.A_GEN_WAV_UP_GAMMA, ""),
|
||||
}
|
||||
for bid in range(len(wav_config["upsampling_ratios"])):
|
||||
for key, (tensor_id, suffix) in up_map.items():
|
||||
yield (self.format_tensor_name(tensor_id, bid, suffix=suffix), get(f"decoder.upsample.{bid}.{key}"))
|
||||
|
||||
# --- DAC decoder ---
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_ENTRY, suffix=".weight"), get("decoder.decoder.0.conv.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_ENTRY, suffix=".bias"), get("decoder.decoder.0.conv.bias"))
|
||||
|
||||
n_dac_blocks = len(wav_config["upsample_rates"])
|
||||
for bid in range(n_dac_blocks):
|
||||
py = bid + 1 # decoder.decoder.0 is the entry conv, blocks start at 1
|
||||
|
||||
a, b = snake_fold(get(f"decoder.decoder.{py}.block.0.alpha"), get(f"decoder.decoder.{py}.block.0.beta"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_SNAKE, bid, suffix=".alpha"), a)
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_SNAKE, bid, suffix=".beta"), b)
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_CONV, bid, suffix=".weight"), get(f"decoder.decoder.{py}.block.1.conv.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_CONV, bid, suffix=".bias"), get(f"decoder.decoder.{py}.block.1.conv.bias"))
|
||||
|
||||
for xid in range(3):
|
||||
ridx = xid + 2 # block.2/3/4 are the 3 residual units
|
||||
|
||||
a1, b1 = snake_fold(get(f"decoder.decoder.{py}.block.{ridx}.act1.alpha"), get(f"decoder.decoder.{py}.block.{ridx}.act1.beta"))
|
||||
name1 = gguf.TENSOR_NAMES[T.A_GEN_WAV_DAC_RES_ACT1].format(bid=bid, xid=xid)
|
||||
yield (name1 + ".alpha", a1)
|
||||
yield (name1 + ".beta", b1)
|
||||
|
||||
name_conv1 = gguf.TENSOR_NAMES[T.A_GEN_WAV_DAC_RES_CONV1].format(bid=bid, xid=xid)
|
||||
yield (name_conv1 + ".weight", get(f"decoder.decoder.{py}.block.{ridx}.conv1.conv.weight"))
|
||||
yield (name_conv1 + ".bias", get(f"decoder.decoder.{py}.block.{ridx}.conv1.conv.bias"))
|
||||
|
||||
a2, b2 = snake_fold(get(f"decoder.decoder.{py}.block.{ridx}.act2.alpha"), get(f"decoder.decoder.{py}.block.{ridx}.act2.beta"))
|
||||
name2 = gguf.TENSOR_NAMES[T.A_GEN_WAV_DAC_RES_ACT2].format(bid=bid, xid=xid)
|
||||
yield (name2 + ".alpha", a2)
|
||||
yield (name2 + ".beta", b2)
|
||||
|
||||
name_conv2 = gguf.TENSOR_NAMES[T.A_GEN_WAV_DAC_RES_CONV2].format(bid=bid, xid=xid)
|
||||
yield (name_conv2 + ".weight", get(f"decoder.decoder.{py}.block.{ridx}.conv2.conv.weight"))
|
||||
yield (name_conv2 + ".bias", get(f"decoder.decoder.{py}.block.{ridx}.conv2.conv.bias"))
|
||||
|
||||
a5, b5 = snake_fold(get("decoder.decoder.5.alpha"), get("decoder.decoder.5.beta"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_SNAKE, suffix=".alpha"), a5)
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_SNAKE, suffix=".beta"), b5)
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_CONV, suffix=".weight"), get("decoder.decoder.6.conv.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_CONV, suffix=".bias"), get("decoder.decoder.6.conv.bias"))
|
||||
@@ -133,6 +133,7 @@ Note:
|
||||
- To debug the multimodal preprocessor and encoder, you can use [llama-mtmd-debug](tools/mtmd/debug/mtmd-debug.cpp).
|
||||
- Adding a model-specific API or CLI is an anti-pattern in `libmtmd`. The goal of `libmtmd` is to provide an easy-to-use, model-agnostic library for multimodal pipeline.
|
||||
- In most cases, `llama-mtmd-cli` should not be modified. If a model requires a specific prompt, either let the user provide it or bake it into the Jinja chat template.
|
||||
- For audio generation models, see `tools/mtmd/README-dev.md`
|
||||
|
||||
## Tips and tricks
|
||||
|
||||
|
||||
@@ -341,6 +341,9 @@ set(GGML_PUBLIC_HEADERS
|
||||
include/gguf.h)
|
||||
|
||||
set_target_properties(ggml PROPERTIES PUBLIC_HEADER "${GGML_PUBLIC_HEADERS}")
|
||||
#if (GGML_METAL)
|
||||
# set_target_properties(ggml PROPERTIES RESOURCE "${CMAKE_CURRENT_SOURCE_DIR}/src/ggml-metal.metal")
|
||||
#endif()
|
||||
install(TARGETS ggml LIBRARY PUBLIC_HEADER)
|
||||
install(TARGETS ggml-base LIBRARY)
|
||||
|
||||
|
||||
@@ -24,119 +24,62 @@ if (GGML_METAL_NDEBUG)
|
||||
endif()
|
||||
|
||||
set(METALLIB_COMMON "${CMAKE_CURRENT_SOURCE_DIR}/../ggml-common.h")
|
||||
set(METALLIB_KERNELS_COMMON "${CMAKE_CURRENT_SOURCE_DIR}/kernels/common.h")
|
||||
set(METALLIB_KERNELS_DEQUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/dequantize.h")
|
||||
set(METALLIB_KERNELS_QUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/quantize.h")
|
||||
|
||||
set(METALLIB_KERNEL_SOURCES
|
||||
kernels/fa.metal
|
||||
kernels/mul_mv.metal
|
||||
kernels/mul_mm.metal
|
||||
kernels/quantize.metal
|
||||
kernels/softmax.metal
|
||||
kernels/norm.metal
|
||||
kernels/unary.metal
|
||||
kernels/binbcast.metal
|
||||
kernels/reduce.metal
|
||||
kernels/tri.metal
|
||||
kernels/ssm.metal
|
||||
kernels/wkv.metal
|
||||
kernels/gated_delta_net.metal
|
||||
kernels/solve_tri.metal
|
||||
kernels/rope.metal
|
||||
kernels/conv.metal
|
||||
kernels/upscale.metal
|
||||
kernels/argsort.metal
|
||||
kernels/pool.metal
|
||||
kernels/misc.metal
|
||||
)
|
||||
|
||||
if (GGML_METAL_EMBED_LIBRARY)
|
||||
enable_language(ASM)
|
||||
|
||||
add_compile_definitions(GGML_METAL_EMBED_LIBRARY)
|
||||
|
||||
set(METALLIB_IMPL "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal-impl.h")
|
||||
set(METALLIB_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal.metal")
|
||||
set(METALLIB_IMPL "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal-impl.h")
|
||||
|
||||
file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/autogenerated")
|
||||
|
||||
set(METALLIB_EMBED_ASM_FILES "")
|
||||
foreach(src ${METALLIB_KERNEL_SOURCES})
|
||||
get_filename_component(kind ${src} NAME_WE)
|
||||
# symbol names must be valid C identifiers ('-' is not allowed)
|
||||
string(REPLACE "-" "_" kind_sym ${kind})
|
||||
# merge ggml-common.h and ggml-metal.metal into a single file
|
||||
set(METALLIB_EMBED_ASM "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.s")
|
||||
set(METALLIB_SOURCE_EMBED "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.metal")
|
||||
set(METALLIB_SOURCE_EMBED_TMP "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.metal.tmp")
|
||||
|
||||
set(SRC "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${kind}.metal")
|
||||
set(EMBED "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed-${kind}.metal")
|
||||
set(ASM "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed-${kind}.s")
|
||||
add_custom_command(
|
||||
OUTPUT "${METALLIB_EMBED_ASM}"
|
||||
COMMAND echo "Embedding Metal library"
|
||||
COMMAND sed -e "/__embed_ggml-common.h__/r ${METALLIB_COMMON}" -e "/__embed_ggml-common.h__/d" < "${METALLIB_SOURCE}" > "${METALLIB_SOURCE_EMBED_TMP}"
|
||||
COMMAND sed -e "/\#include \"ggml-metal-impl.h\"/r ${METALLIB_IMPL}" -e "/\#include \"ggml-metal-impl.h\"/d" < "${METALLIB_SOURCE_EMBED_TMP}" > "${METALLIB_SOURCE_EMBED}"
|
||||
COMMAND echo ".section __DATA,__ggml_metallib" > "${METALLIB_EMBED_ASM}"
|
||||
COMMAND echo ".globl _ggml_metallib_start" >> "${METALLIB_EMBED_ASM}"
|
||||
COMMAND echo "_ggml_metallib_start:" >> "${METALLIB_EMBED_ASM}"
|
||||
COMMAND echo .incbin "\"${METALLIB_SOURCE_EMBED}\"" >> "${METALLIB_EMBED_ASM}"
|
||||
COMMAND echo ".globl _ggml_metallib_end" >> "${METALLIB_EMBED_ASM}"
|
||||
COMMAND echo "_ggml_metallib_end:" >> "${METALLIB_EMBED_ASM}"
|
||||
DEPENDS ../ggml-common.h ggml-metal.metal ggml-metal-impl.h
|
||||
COMMENT "Generate assembly for embedded Metal library"
|
||||
VERBATIM
|
||||
)
|
||||
|
||||
# only prepend headers that this source actually includes
|
||||
set(HEADERS_FOR_SRC ${METALLIB_KERNELS_COMMON})
|
||||
file(STRINGS ${SRC} _has_dequantize REGEX "#include \"dequantize\\.h\"")
|
||||
file(STRINGS ${SRC} _has_quantize REGEX "#include \"quantize\\.h\"")
|
||||
if(_has_dequantize)
|
||||
list(APPEND HEADERS_FOR_SRC ${METALLIB_KERNELS_DEQUANTIZE})
|
||||
endif()
|
||||
if(_has_quantize)
|
||||
list(APPEND HEADERS_FOR_SRC ${METALLIB_KERNELS_QUANTIZE})
|
||||
endif()
|
||||
|
||||
add_custom_command(
|
||||
OUTPUT "${ASM}"
|
||||
# Step 1: concatenate shared headers + this kernel source
|
||||
COMMAND cat ${HEADERS_FOR_SRC} ${SRC} > "${EMBED}.tmp1"
|
||||
# Step 2: remove internal #include and #pragma once
|
||||
COMMAND sed -e "/\#include \"common.h\"/d" -e "/\#include \"dequantize.h\"/d" -e "/\#include \"quantize.h\"/d" -e "/\#pragma once/d" < "${EMBED}.tmp1" > "${EMBED}.tmp2"
|
||||
# Step 3: inline ggml-common.h (replacing __embed_ggml-common.h__ sentinel)
|
||||
COMMAND sed -e "/__embed_ggml-common.h__/r ${METALLIB_COMMON}" -e "/__embed_ggml-common.h__/d" < "${EMBED}.tmp2" > "${EMBED}.tmp3"
|
||||
# Step 4: inline ggml-metal-impl.h
|
||||
COMMAND sed -e "/\#include \"ggml-metal-impl.h\"/r ${METALLIB_IMPL}" -e "/\#include \"ggml-metal-impl.h\"/d" < "${EMBED}.tmp3" > "${EMBED}"
|
||||
# Step 5: emit an asm chunk with kind-specific start/end symbols
|
||||
# note: '-' is illegal in C symbols, so we use kind_sym; the macOS
|
||||
# section name is limited to 16 chars so we keep it shared
|
||||
# across kinds (__ggml_metallib) and only vary the global symbols.
|
||||
COMMAND echo ".section __DATA,__ggml_metallib" > "${ASM}"
|
||||
COMMAND echo ".globl _ggml_metallib_${kind_sym}_start" >> "${ASM}"
|
||||
COMMAND echo "_ggml_metallib_${kind_sym}_start:" >> "${ASM}"
|
||||
COMMAND echo .incbin "\"${EMBED}\"" >> "${ASM}"
|
||||
COMMAND echo ".globl _ggml_metallib_${kind_sym}_end" >> "${ASM}"
|
||||
COMMAND echo "_ggml_metallib_${kind_sym}_end:" >> "${ASM}"
|
||||
DEPENDS ../ggml-common.h ggml-metal-impl.h
|
||||
kernels/common.h kernels/dequantize.h kernels/quantize.h
|
||||
kernels/${kind}.metal
|
||||
COMMENT "Generate embedded Metal library for ${kind}"
|
||||
VERBATIM
|
||||
)
|
||||
|
||||
list(APPEND METALLIB_EMBED_ASM_FILES "${ASM}")
|
||||
endforeach()
|
||||
|
||||
target_sources(ggml-metal PRIVATE ${METALLIB_EMBED_ASM_FILES})
|
||||
target_sources(ggml-metal PRIVATE "${METALLIB_EMBED_ASM}")
|
||||
else()
|
||||
# copy header files to bin directory
|
||||
# copy metal files to bin directory
|
||||
configure_file(../ggml-common.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h COPYONLY)
|
||||
configure_file(ggml-metal.metal ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal COPYONLY)
|
||||
configure_file(ggml-metal-impl.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h COPYONLY)
|
||||
|
||||
file(MAKE_DIRECTORY "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels")
|
||||
configure_file(kernels/common.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/common.h COPYONLY)
|
||||
configure_file(kernels/dequantize.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/dequantize.h COPYONLY)
|
||||
configure_file(kernels/quantize.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/quantize.h COPYONLY)
|
||||
|
||||
foreach(src ${METALLIB_KERNEL_SOURCES})
|
||||
configure_file(${src} ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} COPYONLY)
|
||||
endforeach()
|
||||
|
||||
if (GGML_METAL_SHADER_DEBUG)
|
||||
# note: disabling fast math is needed in order to pass tests/test-backend-ops
|
||||
# custom command to do the following:
|
||||
# xcrun -sdk macosx metal -fno-fast-math -c ggml-metal.metal -o ggml-metal.air
|
||||
# xcrun -sdk macosx metallib ggml-metal.air -o default.metallib
|
||||
#
|
||||
# note: this is the only way I found to disable fast-math in Metal. it's ugly, but at least it works
|
||||
# disabling fast math is needed in order to pass tests/test-backend-ops
|
||||
# note: adding -fno-inline fixes the tests when using MTL_SHADER_VALIDATION=1
|
||||
# note: unfortunately, we have to call it default.metallib instead of ggml.metallib
|
||||
# ref: https://github.com/ggml-org/whisper.cpp/issues/1720
|
||||
# note: adding -g causes segmentation fault during compile
|
||||
#set(XC_FLAGS -fno-fast-math -fno-inline -g)
|
||||
set(XC_FLAGS -fno-fast-math -fno-inline)
|
||||
else()
|
||||
set(XC_FLAGS -O3)
|
||||
endif()
|
||||
|
||||
# Append macOS metal versioning flags
|
||||
if (GGML_METAL_MACOSX_VERSION_MIN)
|
||||
message(STATUS "Adding -mmacosx-version-min=${GGML_METAL_MACOSX_VERSION_MIN} flag to metal compilation")
|
||||
list (APPEND XC_FLAGS -mmacosx-version-min=${GGML_METAL_MACOSX_VERSION_MIN})
|
||||
@@ -147,46 +90,35 @@ else()
|
||||
list (APPEND XC_FLAGS -std=${GGML_METAL_STD})
|
||||
endif()
|
||||
|
||||
# Compile each kernel source to .air, then link into default.metallib
|
||||
set(AIR_FILES "")
|
||||
foreach(src ${METALLIB_KERNEL_SOURCES})
|
||||
get_filename_component(name ${src} NAME_WE)
|
||||
set(AIR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${name}.air")
|
||||
list(APPEND AIR_FILES ${AIR})
|
||||
add_custom_command(
|
||||
OUTPUT ${AIR}
|
||||
COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} -o ${AIR}
|
||||
DEPENDS ${src} kernels/common.h kernels/dequantize.h kernels/quantize.h ${METALLIB_COMMON} ggml-metal-impl.h
|
||||
COMMENT "Compiling ${src}"
|
||||
VERBATIM
|
||||
)
|
||||
endforeach()
|
||||
|
||||
add_custom_command(
|
||||
OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
|
||||
COMMAND xcrun -sdk macosx metallib ${AIR_FILES} -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
|
||||
COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal -o - |
|
||||
xcrun -sdk macosx metallib - -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
|
||||
COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h
|
||||
COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h
|
||||
COMMAND rm -rf ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels
|
||||
DEPENDS ${AIR_FILES}
|
||||
COMMENT "Linking Metal kernels into default.metallib"
|
||||
)
|
||||
COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal
|
||||
DEPENDS ggml-metal.metal ${METALLIB_COMMON}
|
||||
COMMENT "Compiling Metal kernels"
|
||||
)
|
||||
|
||||
# FIXME: only add to the ggml-metal target?
|
||||
add_custom_target(
|
||||
ggml-metal-lib ALL
|
||||
DEPENDS ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
|
||||
)
|
||||
)
|
||||
endif() # GGML_METAL_EMBED_LIBRARY
|
||||
|
||||
if (NOT GGML_METAL_EMBED_LIBRARY)
|
||||
install(
|
||||
DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/kernels/
|
||||
DESTINATION ${CMAKE_INSTALL_BINDIR}/kernels
|
||||
FILES_MATCHING PATTERN "*.metal" PATTERN "*.h"
|
||||
)
|
||||
FILES src/ggml-metal/ggml-metal.metal
|
||||
PERMISSIONS
|
||||
OWNER_READ
|
||||
OWNER_WRITE
|
||||
GROUP_READ
|
||||
WORLD_READ
|
||||
DESTINATION ${CMAKE_INSTALL_BINDIR})
|
||||
|
||||
install(
|
||||
FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
|
||||
DESTINATION ${CMAKE_INSTALL_BINDIR}
|
||||
)
|
||||
install(
|
||||
FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
|
||||
DESTINATION ${CMAKE_INSTALL_BINDIR}
|
||||
)
|
||||
endif()
|
||||
|
||||
@@ -95,63 +95,8 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi
|
||||
return pipeline.pipeline->obj.maxTotalThreadsPerThreadgroup;
|
||||
}
|
||||
|
||||
//
|
||||
// MTLLibrary collection (one library per op-source, compiled separately)
|
||||
//
|
||||
|
||||
// Single source of truth for the per-kind metal libraries. The order here
|
||||
// defines the enum values and every per-kind table below, so adding a library
|
||||
// is a one-line change here (plus adding its source to CMakeLists.txt).
|
||||
// X(suffix, name): name is both the kernels/<name>.metal basename and the
|
||||
// ggml_metallib_<name>_{start,end} embed-symbol stem.
|
||||
#define GGML_METAL_LIBS \
|
||||
X(FA, fa) \
|
||||
X(MUL_MV, mul_mv) \
|
||||
X(MUL_MM, mul_mm) \
|
||||
X(QUANTIZE, quantize) \
|
||||
X(SOFTMAX, softmax) \
|
||||
X(NORM, norm) \
|
||||
X(UNARY, unary) \
|
||||
X(BINBCAST, binbcast) \
|
||||
X(REDUCE, reduce) \
|
||||
X(TRI, tri) \
|
||||
X(SSM, ssm) \
|
||||
X(WKV, wkv) \
|
||||
X(GATED_DELTA_NET, gated_delta_net)\
|
||||
X(SOLVE_TRI, solve_tri) \
|
||||
X(ROPE, rope) \
|
||||
X(CONV, conv) \
|
||||
X(UPSCALE, upscale) \
|
||||
X(ARGSORT, argsort) \
|
||||
X(POOL, pool) \
|
||||
X(MISC, misc)
|
||||
|
||||
enum ggml_metal_lib_kind {
|
||||
#define X(e, s) GGML_METAL_LIB_##e,
|
||||
GGML_METAL_LIBS
|
||||
#undef X
|
||||
GGML_METAL_LIB_COUNT,
|
||||
};
|
||||
|
||||
static const char * const k_lib_names[GGML_METAL_LIB_COUNT] = {
|
||||
#define X(e, s) [GGML_METAL_LIB_##e] = #s,
|
||||
GGML_METAL_LIBS
|
||||
#undef X
|
||||
};
|
||||
|
||||
struct ggml_metal_library {
|
||||
// Per-kind compiled libraries. When single_library is true, the whole library
|
||||
// (e.g. a pre-compiled default.metallib or a from-source build) lives at
|
||||
// objs[0] and the remaining slots are nil.
|
||||
id<MTLLibrary> objs[GGML_METAL_LIB_COUNT];
|
||||
bool single_library; // true: combined library at objs[0]; false: per-kind libs in objs[*]
|
||||
|
||||
// Routing table: kernel function name -> objs[] index, populated from each
|
||||
// compiled library's -[MTLLibrary functionNames]. The actual compiled
|
||||
// libraries are the single source of truth for which library owns a kernel,
|
||||
// so adding kernels later requires no manual routing maintenance.
|
||||
// nil in single_library mode (everything resolves to objs[0]).
|
||||
NSMutableDictionary<NSString *, NSNumber *> * fn_to_lib;
|
||||
id<MTLLibrary> obj;
|
||||
|
||||
ggml_metal_device_t dev;
|
||||
ggml_metal_pipelines_t pipelines; // cache of compiled pipelines
|
||||
@@ -159,376 +104,160 @@ struct ggml_metal_library {
|
||||
NSLock * lock;
|
||||
};
|
||||
|
||||
// Build the fn_to_lib routing table by querying each compiled library's public
|
||||
// function names. Call once after all per-kind libraries have been compiled.
|
||||
static void ggml_metal_library_build_index(ggml_metal_library_t lib) {
|
||||
@autoreleasepool {
|
||||
NSMutableDictionary<NSString *, NSNumber *> * index = [[NSMutableDictionary alloc] init];
|
||||
for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
|
||||
for (NSString * fname in [lib->objs[kind] functionNames]) {
|
||||
index[fname] = @(kind);
|
||||
}
|
||||
}
|
||||
lib->fn_to_lib = index;
|
||||
}
|
||||
}
|
||||
ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
|
||||
id<MTLLibrary> library = nil;
|
||||
id<MTLDevice> device = ggml_metal_device_get_obj(dev);
|
||||
|
||||
// Parse a `#include "name"` line. Returns the quoted name in *include_name on
|
||||
// success. Whitespace-tolerant; ignores `#include <...>` (system headers).
|
||||
static bool ggml_metal_library_parse_quoted_include(NSString * line, NSString ** include_name) {
|
||||
NSScanner * scanner = [NSScanner scannerWithString:line];
|
||||
scanner.charactersToBeSkipped = [NSCharacterSet whitespaceCharacterSet];
|
||||
// load library
|
||||
//
|
||||
// - first check if the library is embedded
|
||||
// - then check if the library is in the bundle
|
||||
// - if not found, load the source and compile it
|
||||
// - if that fails, return NULL
|
||||
//
|
||||
// TODO: move to a function
|
||||
{
|
||||
const int64_t t_start = ggml_time_us();
|
||||
|
||||
if (![scanner scanString:@"#" intoString:NULL] ||
|
||||
![scanner scanString:@"include" intoString:NULL] ||
|
||||
![scanner scanString:@"\"" intoString:NULL]) {
|
||||
return false;
|
||||
}
|
||||
NSError * error = nil;
|
||||
NSString * src = nil;
|
||||
|
||||
NSString * name = nil;
|
||||
if (![scanner scanUpToString:@"\"" intoString:&name]) {
|
||||
return false;
|
||||
}
|
||||
#if GGML_METAL_EMBED_LIBRARY
|
||||
GGML_LOG_INFO("%s: using embedded metal library\n", __func__);
|
||||
|
||||
if (include_name) {
|
||||
*include_name = name;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
extern const char ggml_metallib_start[];
|
||||
extern const char ggml_metallib_end[];
|
||||
|
||||
// Recursively inline `#include "name"` directives. System includes (<...>),
|
||||
// `#if/#else/#endif`, and other preprocessor lines are passed through to the
|
||||
// Metal compiler unchanged. `#pragma once` is dropped since `seen` already
|
||||
// guards against double-inclusion.
|
||||
static bool ggml_metal_library_flatten_file(NSMutableString * dst, NSString * path,
|
||||
NSArray<NSString *> * search_paths,
|
||||
NSMutableSet<NSString *> * seen, NSError ** error) {
|
||||
NSString * key = [path stringByStandardizingPath];
|
||||
if ([seen containsObject:key]) {
|
||||
return true;
|
||||
}
|
||||
[seen addObject:key];
|
||||
src = [[NSString alloc] initWithBytes:ggml_metallib_start length:(ggml_metallib_end-ggml_metallib_start) encoding:NSUTF8StringEncoding];
|
||||
#else
|
||||
|
||||
NSString * src = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:error];
|
||||
if (!src) {
|
||||
return false;
|
||||
}
|
||||
#ifdef SWIFT_PACKAGE
|
||||
NSBundle * bundle = SWIFTPM_MODULE_BUNDLE;
|
||||
#else
|
||||
NSBundle * bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
|
||||
#endif
|
||||
|
||||
NSFileManager * fm = [NSFileManager defaultManager];
|
||||
for (NSString * line in [src componentsSeparatedByString:@"\n"]) {
|
||||
NSString * trimmed = [line stringByTrimmingCharactersInSet:[NSCharacterSet whitespaceCharacterSet]];
|
||||
if ([trimmed isEqualToString:@"#pragma once"]) {
|
||||
continue;
|
||||
}
|
||||
NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
|
||||
if (path_lib == nil) {
|
||||
// Try to find the resource in the directory where the current binary located.
|
||||
NSString * bin_cur = [[NSProcessInfo processInfo] arguments][0];
|
||||
NSString * bin_dir = [bin_cur stringByDeletingLastPathComponent];
|
||||
|
||||
NSString * include_name = nil;
|
||||
if (ggml_metal_library_parse_quoted_include(line, &include_name)) {
|
||||
NSString * resolved = nil;
|
||||
for (NSString * dir in search_paths) {
|
||||
NSString * candidate = [dir stringByAppendingPathComponent:include_name];
|
||||
if ([fm isReadableFileAtPath:candidate]) {
|
||||
resolved = candidate;
|
||||
break;
|
||||
NSString * path_lib_default = [NSString pathWithComponents:@[bin_dir, @"default.metallib"]];
|
||||
if ([[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
|
||||
GGML_LOG_INFO("%s: found '%s'\n", __func__, [path_lib_default UTF8String]);
|
||||
|
||||
NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:path_lib_default error:&error];
|
||||
if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
|
||||
// Optionally, if this is a symlink, try to resolve it.
|
||||
path_lib_default = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:path_lib_default error:&error];
|
||||
if (path_lib_default && [path_lib_default length] > 0 && ![[path_lib_default substringToIndex:1] isEqualToString:@"/"]) {
|
||||
// It is a relative path, adding the binary directory as directory prefix.
|
||||
path_lib_default = [NSString pathWithComponents:@[bin_dir, path_lib_default]];
|
||||
}
|
||||
if (!path_lib_default || ![[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
|
||||
// Link to the resource could not be resolved.
|
||||
path_lib_default = nil;
|
||||
} else {
|
||||
GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [path_lib_default UTF8String]);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// The resource couldn't be found in the binary's directory.
|
||||
path_lib_default = nil;
|
||||
}
|
||||
if (!resolved) {
|
||||
if (error) {
|
||||
NSString * msg = [NSString stringWithFormat:@"could not resolve include \"%@\" from '%@'", include_name, path];
|
||||
*error = [NSError errorWithDomain:@"ggml-metal-source-flatten" code:1
|
||||
userInfo:@{NSLocalizedDescriptionKey: msg}];
|
||||
}
|
||||
return false;
|
||||
}
|
||||
if (!ggml_metal_library_flatten_file(dst, resolved, search_paths, seen, error)) {
|
||||
return false;
|
||||
}
|
||||
continue;
|
||||
|
||||
path_lib = path_lib_default;
|
||||
}
|
||||
|
||||
[dst appendString:line];
|
||||
[dst appendString:@"\n"];
|
||||
}
|
||||
if (path_lib != nil) {
|
||||
// pre-compiled library found
|
||||
NSURL * libURL = [NSURL fileURLWithPath:path_lib];
|
||||
GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_lib UTF8String]);
|
||||
|
||||
return true;
|
||||
}
|
||||
library = [device newLibraryWithURL:libURL error:&error];
|
||||
if (error) {
|
||||
GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
|
||||
return nil;
|
||||
}
|
||||
} else {
|
||||
GGML_LOG_INFO("%s: default.metallib not found, loading from source\n", __func__);
|
||||
|
||||
static NSString * ggml_metal_library_flatten_source(NSString * path_source, NSError ** error) {
|
||||
// Search paths cover both runtime layout (build/bin/kernels + build/bin)
|
||||
// and source-tree layout (ggml/src/ggml-metal/kernels + ggml/src/ggml-metal + ggml/src).
|
||||
NSString * path_kernels = [path_source stringByDeletingLastPathComponent];
|
||||
NSString * path_base = [path_kernels stringByDeletingLastPathComponent];
|
||||
NSArray<NSString *> * search_paths = @[
|
||||
path_kernels,
|
||||
path_base,
|
||||
[path_base stringByDeletingLastPathComponent],
|
||||
];
|
||||
NSString * path_source;
|
||||
NSString * path_resource = [[NSProcessInfo processInfo].environment objectForKey:@"GGML_METAL_PATH_RESOURCES"];
|
||||
|
||||
NSMutableString * src = [[NSMutableString alloc] init];
|
||||
NSMutableSet<NSString *> * seen = [NSMutableSet set];
|
||||
GGML_LOG_INFO("%s: GGML_METAL_PATH_RESOURCES = %s\n", __func__, path_resource ? [path_resource UTF8String] : "nil");
|
||||
|
||||
if (!ggml_metal_library_flatten_file(src, path_source, search_paths, seen, error)) {
|
||||
[src release];
|
||||
return nil;
|
||||
}
|
||||
return src;
|
||||
}
|
||||
|
||||
// Compile all per-kind libraries in parallel. `source_for_kind` returns the MSL
|
||||
// source for a kind (the helper takes ownership and releases it), or nil with
|
||||
// *err set on failure. On success the objs[] slots are populated and the routing
|
||||
// index is built; on any failure every error is logged and false is returned
|
||||
// (the caller is responsible for freeing `res`).
|
||||
static bool ggml_metal_library_compile_all(
|
||||
ggml_metal_library_t res,
|
||||
id<MTLDevice> device,
|
||||
NSDictionary * prep,
|
||||
NSString * (^source_for_kind)(int kind, NSError ** err),
|
||||
const char * origin) {
|
||||
const int64_t t_start = ggml_time_us();
|
||||
|
||||
int64_t * t_per_lib = calloc(GGML_METAL_LIB_COUNT, sizeof(int64_t));
|
||||
NSError ** err_per_lib = calloc(GGML_METAL_LIB_COUNT, sizeof(NSError *));
|
||||
__block atomic_bool any_failure = false;
|
||||
|
||||
dispatch_group_t group = dispatch_group_create();
|
||||
dispatch_queue_t queue = dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0);
|
||||
|
||||
for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
|
||||
dispatch_group_async(group, queue, ^{
|
||||
|
||||
const int64_t t0 = ggml_time_us();
|
||||
|
||||
NSError * error = nil;
|
||||
|
||||
NSString * src = source_for_kind(kind, &error);
|
||||
if (!src) {
|
||||
err_per_lib[kind] = [error retain];
|
||||
atomic_store(&any_failure, true);
|
||||
return;
|
||||
if (path_resource) {
|
||||
path_source = [path_resource stringByAppendingPathComponent:@"ggml-metal.metal"];
|
||||
} else {
|
||||
path_source = [bundle pathForResource:@"ggml-metal" ofType:@"metal"];
|
||||
}
|
||||
|
||||
id<MTLLibrary> lib = nil;
|
||||
if (path_source == nil) {
|
||||
GGML_LOG_WARN("%s: error: could not use bundle path to find ggml-metal.metal, falling back to trying cwd\n", __func__);
|
||||
path_source = @"ggml-metal.metal";
|
||||
}
|
||||
|
||||
GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_source UTF8String]);
|
||||
|
||||
src = [NSString stringWithContentsOfFile:path_source encoding:NSUTF8StringEncoding error:&error];
|
||||
if (error) {
|
||||
GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
|
||||
return nil;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
if (!library) {
|
||||
@autoreleasepool {
|
||||
// dictionary of preprocessor macros
|
||||
NSMutableDictionary * prep = [NSMutableDictionary dictionary];
|
||||
|
||||
if (ggml_metal_device_get_props(dev)->has_bfloat) {
|
||||
[prep setObject:@"1" forKey:@"GGML_METAL_HAS_BF16"];
|
||||
}
|
||||
|
||||
if (ggml_metal_device_get_props(dev)->has_tensor) {
|
||||
[prep setObject:@"1" forKey:@"GGML_METAL_HAS_TENSOR"];
|
||||
}
|
||||
|
||||
#if GGML_METAL_EMBED_LIBRARY
|
||||
[prep setObject:@"1" forKey:@"GGML_METAL_EMBED_LIBRARY"];
|
||||
#endif
|
||||
|
||||
MTLCompileOptions * options = [MTLCompileOptions new];
|
||||
options.preprocessorMacros = prep;
|
||||
|
||||
lib = [device newLibraryWithSource:src options:options error:&error];
|
||||
//[options setFastMathEnabled:false];
|
||||
|
||||
[options release];
|
||||
|
||||
// retain the error before the autorelease pool drains it
|
||||
if (!lib) {
|
||||
err_per_lib[kind] = [error retain];
|
||||
library = [device newLibraryWithSource:src options:options error:&error];
|
||||
if (error) {
|
||||
GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
|
||||
return nil;
|
||||
}
|
||||
}
|
||||
|
||||
[src release];
|
||||
|
||||
t_per_lib[kind] = ggml_time_us() - t0;
|
||||
|
||||
if (!lib) {
|
||||
atomic_store(&any_failure, true);
|
||||
return;
|
||||
}
|
||||
|
||||
res->objs[kind] = lib;
|
||||
});
|
||||
}
|
||||
dispatch_group_wait(group, DISPATCH_TIME_FOREVER);
|
||||
dispatch_release(group);
|
||||
|
||||
const bool ok = !atomic_load(&any_failure);
|
||||
|
||||
if (ok) {
|
||||
const int64_t t_total = ggml_time_us() - t_start;
|
||||
int64_t t_max = 0;
|
||||
for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
|
||||
GGML_LOG_DEBUG("%s: compiled '%s' library in %.3f sec\n",
|
||||
__func__, k_lib_names[kind], t_per_lib[kind] / 1e6);
|
||||
if (t_per_lib[kind] > t_max) t_max = t_per_lib[kind];
|
||||
}
|
||||
GGML_LOG_INFO("%s: loaded %d libraries from %s in %.3f sec (max single = %.3f sec)\n",
|
||||
__func__, GGML_METAL_LIB_COUNT, origin, t_total / 1e6, t_max / 1e6);
|
||||
|
||||
ggml_metal_library_build_index(res);
|
||||
} else {
|
||||
for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
|
||||
if (err_per_lib[kind]) {
|
||||
GGML_LOG_ERROR("%s: failed to build '%s' library: %s\n", __func__,
|
||||
k_lib_names[kind], [[err_per_lib[kind] description] UTF8String]);
|
||||
[err_per_lib[kind] release];
|
||||
#if !__has_feature(objc_arc)
|
||||
[options release];
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
#if GGML_METAL_EMBED_LIBRARY
|
||||
[src release];
|
||||
#endif // GGML_METAL_EMBED_LIBRARY
|
||||
|
||||
GGML_LOG_INFO("%s: loaded in %.3f sec\n", __func__, (ggml_time_us() - t_start) / 1e6);
|
||||
}
|
||||
|
||||
free(err_per_lib);
|
||||
free(t_per_lib);
|
||||
|
||||
return ok;
|
||||
}
|
||||
|
||||
ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
|
||||
id<MTLDevice> device = ggml_metal_device_get_obj(dev);
|
||||
|
||||
ggml_metal_library_t res = calloc(1, sizeof(struct ggml_metal_library));
|
||||
|
||||
res->obj = library;
|
||||
res->dev = dev;
|
||||
res->pipelines = ggml_metal_pipelines_init();
|
||||
res->lock = [NSLock new];
|
||||
|
||||
// shared MTLCompileOptions preprocessor macros (matches the build-time defines)
|
||||
NSMutableDictionary * prep = [NSMutableDictionary dictionary];
|
||||
if (ggml_metal_device_get_props(dev)->has_bfloat) {
|
||||
[prep setObject:@"1" forKey:@"GGML_METAL_HAS_BF16"];
|
||||
}
|
||||
if (ggml_metal_device_get_props(dev)->has_tensor) {
|
||||
[prep setObject:@"1" forKey:@"GGML_METAL_HAS_TENSOR"];
|
||||
}
|
||||
#if GGML_METAL_EMBED_LIBRARY
|
||||
[prep setObject:@"1" forKey:@"GGML_METAL_EMBED_LIBRARY"];
|
||||
#endif
|
||||
|
||||
#if GGML_METAL_EMBED_LIBRARY
|
||||
GGML_LOG_INFO("%s: using embedded metal library\n", __func__);
|
||||
|
||||
// start/end symbols emitted by CMake (see CMakeLists.txt), one pair per kind
|
||||
#define X(e, s) extern const char ggml_metallib_##s##_start[]; extern const char ggml_metallib_##s##_end[];
|
||||
GGML_METAL_LIBS
|
||||
#undef X
|
||||
|
||||
static const char * const lib_start[GGML_METAL_LIB_COUNT] = {
|
||||
#define X(e, s) [GGML_METAL_LIB_##e] = ggml_metallib_##s##_start,
|
||||
GGML_METAL_LIBS
|
||||
#undef X
|
||||
};
|
||||
static const char * const lib_end[GGML_METAL_LIB_COUNT] = {
|
||||
#define X(e, s) [GGML_METAL_LIB_##e] = ggml_metallib_##s##_end,
|
||||
GGML_METAL_LIBS
|
||||
#undef X
|
||||
};
|
||||
|
||||
const bool ok = ggml_metal_library_compile_all(res, device, prep,
|
||||
^NSString * (int kind, NSError ** err) {
|
||||
(void) err;
|
||||
return [[NSString alloc] initWithBytes:lib_start[kind]
|
||||
length:(lib_end[kind] - lib_start[kind])
|
||||
encoding:NSUTF8StringEncoding];
|
||||
}, "embedded data");
|
||||
|
||||
if (!ok) {
|
||||
ggml_metal_library_free(res);
|
||||
return NULL;
|
||||
}
|
||||
|
||||
return res;
|
||||
#else
|
||||
#ifdef SWIFT_PACKAGE
|
||||
NSBundle * bundle = SWIFTPM_MODULE_BUNDLE;
|
||||
#else
|
||||
NSBundle * bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
|
||||
#endif
|
||||
|
||||
const int64_t t_start = ggml_time_us();
|
||||
|
||||
NSError * error = nil;
|
||||
NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
|
||||
if (path_lib == nil) {
|
||||
// Try to find the resource in the directory where the current binary located.
|
||||
NSString * bin_cur = [[NSProcessInfo processInfo] arguments][0];
|
||||
NSString * bin_dir = [bin_cur stringByDeletingLastPathComponent];
|
||||
|
||||
NSString * path_lib_default = [NSString pathWithComponents:@[bin_dir, @"default.metallib"]];
|
||||
if ([[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
|
||||
GGML_LOG_INFO("%s: found '%s'\n", __func__, [path_lib_default UTF8String]);
|
||||
|
||||
NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:path_lib_default error:&error];
|
||||
if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
|
||||
// Optionally, if this is a symlink, try to resolve it.
|
||||
path_lib_default = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:path_lib_default error:&error];
|
||||
if (path_lib_default && [path_lib_default length] > 0 && ![[path_lib_default substringToIndex:1] isEqualToString:@"/"]) {
|
||||
// It is a relative path, adding the binary directory as directory prefix.
|
||||
path_lib_default = [NSString pathWithComponents:@[bin_dir, path_lib_default]];
|
||||
}
|
||||
if (!path_lib_default || ![[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
|
||||
// Link to the resource could not be resolved.
|
||||
path_lib_default = nil;
|
||||
} else {
|
||||
GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [path_lib_default UTF8String]);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// The resource couldn't be found in the binary's directory.
|
||||
path_lib_default = nil;
|
||||
}
|
||||
|
||||
path_lib = path_lib_default;
|
||||
}
|
||||
|
||||
if (path_lib != nil) {
|
||||
// pre-compiled library found: a single combined default.metallib
|
||||
NSURL * libURL = [NSURL fileURLWithPath:path_lib];
|
||||
GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_lib UTF8String]);
|
||||
|
||||
res->objs[0] = [device newLibraryWithURL:libURL error:&error];
|
||||
res->single_library = true;
|
||||
if (!res->objs[0]) {
|
||||
GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
|
||||
ggml_metal_library_free(res);
|
||||
return NULL;
|
||||
}
|
||||
|
||||
GGML_LOG_INFO("%s: loaded in %.3f sec\n", __func__, (ggml_time_us() - t_start) / 1e6);
|
||||
return res;
|
||||
}
|
||||
|
||||
// no pre-compiled metallib: fall back to compiling each kernel source separately
|
||||
GGML_LOG_INFO("%s: default.metallib not found, loading kernel sources\n", __func__);
|
||||
|
||||
NSString * path_resource = [[NSProcessInfo processInfo].environment objectForKey:@"GGML_METAL_PATH_RESOURCES"];
|
||||
if (path_resource) {
|
||||
GGML_LOG_INFO("%s: GGML_METAL_PATH_RESOURCES = %s\n", __func__, [path_resource UTF8String]);
|
||||
}
|
||||
|
||||
// resolve each kind's source path up front (file lookup/logging stays on the calling thread)
|
||||
NSString ** path_per_kind = calloc(GGML_METAL_LIB_COUNT, sizeof(NSString *));
|
||||
for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
|
||||
NSString * rel = [NSString stringWithFormat:@"kernels/%s.metal", k_lib_names[kind]];
|
||||
|
||||
NSString * path_source = nil;
|
||||
if (path_resource) {
|
||||
path_source = [path_resource stringByAppendingPathComponent:rel];
|
||||
} else {
|
||||
NSString * stem = [NSString stringWithFormat:@"kernels/%s", k_lib_names[kind]];
|
||||
path_source = [bundle pathForResource:stem ofType:@"metal"];
|
||||
}
|
||||
|
||||
if (path_source == nil || ![[NSFileManager defaultManager] isReadableFileAtPath:path_source]) {
|
||||
GGML_LOG_WARN("%s: could not locate %s in bundle, falling back to cwd\n", __func__, [rel UTF8String]);
|
||||
path_source = rel;
|
||||
}
|
||||
|
||||
GGML_LOG_DEBUG("%s: loading '%s'\n", __func__, [path_source UTF8String]);
|
||||
|
||||
path_per_kind[kind] = [path_source retain];
|
||||
}
|
||||
|
||||
const bool ok = ggml_metal_library_compile_all(res, device, prep,
|
||||
^NSString * (int kind, NSError ** err) {
|
||||
return ggml_metal_library_flatten_source(path_per_kind[kind], err);
|
||||
}, "source");
|
||||
|
||||
for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
|
||||
[path_per_kind[kind] release];
|
||||
}
|
||||
free(path_per_kind);
|
||||
|
||||
if (!ok) {
|
||||
ggml_metal_library_free(res);
|
||||
return NULL;
|
||||
}
|
||||
|
||||
return res;
|
||||
#endif
|
||||
}
|
||||
|
||||
ggml_metal_library_t ggml_metal_library_init_from_source(ggml_metal_device_t dev, const char * source, bool verbose) {
|
||||
@@ -590,11 +319,10 @@ ggml_metal_library_t ggml_metal_library_init_from_source(ggml_metal_device_t dev
|
||||
return NULL;
|
||||
}
|
||||
|
||||
res->objs[0] = library;
|
||||
res->single_library = true;
|
||||
res->dev = dev;
|
||||
res->pipelines = ggml_metal_pipelines_init();
|
||||
res->lock = [NSLock new];
|
||||
res->obj = library;
|
||||
res->dev = dev;
|
||||
res->pipelines = ggml_metal_pipelines_init();
|
||||
res->lock = [NSLock new];
|
||||
|
||||
return res;
|
||||
}
|
||||
@@ -604,14 +332,8 @@ void ggml_metal_library_free(ggml_metal_library_t lib) {
|
||||
return;
|
||||
}
|
||||
|
||||
for (int kind = 0; kind < GGML_METAL_LIB_COUNT; ++kind) {
|
||||
if (lib->objs[kind]) {
|
||||
[lib->objs[kind] release];
|
||||
}
|
||||
}
|
||||
|
||||
if (lib->fn_to_lib) {
|
||||
[lib->fn_to_lib release];
|
||||
if (lib->obj) {
|
||||
[lib->obj release];
|
||||
}
|
||||
|
||||
ggml_metal_pipelines_free(lib->pipelines);
|
||||
@@ -672,28 +394,11 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_compile_pipeline(ggml_
|
||||
|
||||
GGML_LOG_DEBUG("%s: compiling pipeline: base = '%s', name = '%s'\n", __func__, base, name);
|
||||
|
||||
// route to the library that actually defines this kernel; fn_to_lib is
|
||||
// built from -[MTLLibrary functionNames] so it's always in sync
|
||||
int lib_idx = 0;
|
||||
if (!lib->single_library) {
|
||||
NSNumber * idx = lib->fn_to_lib[base_func];
|
||||
if (!idx) {
|
||||
[lib->lock unlock];
|
||||
|
||||
GGML_LOG_ERROR("%s: kernel not found in any metal library: base = '%s', name = '%s'\n", __func__, base, name);
|
||||
|
||||
return res;
|
||||
}
|
||||
lib_idx = [idx intValue];
|
||||
}
|
||||
|
||||
id<MTLLibrary> mtl_lib = lib->objs[lib_idx];
|
||||
|
||||
id<MTLFunction> mtl_function;
|
||||
if (!cv) {
|
||||
mtl_function = [mtl_lib newFunctionWithName:base_func];
|
||||
mtl_function = [lib->obj newFunctionWithName:base_func];
|
||||
} else {
|
||||
mtl_function = [mtl_lib newFunctionWithName:base_func constantValues:cv->obj error:&error];
|
||||
mtl_function = [lib->obj newFunctionWithName:base_func constantValues:cv->obj error:&error];
|
||||
}
|
||||
if (!mtl_function) {
|
||||
[lib->lock unlock];
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,232 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
// bitonic sort implementation following the CUDA kernels as reference
|
||||
typedef void (argsort_t)(
|
||||
constant ggml_metal_kargs_argsort & args,
|
||||
device const char * src0,
|
||||
device int32_t * dst,
|
||||
threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
template<ggml_sort_order order>
|
||||
kernel void kernel_argsort_f32_i32(
|
||||
constant ggml_metal_kargs_argsort & args,
|
||||
device const char * src0,
|
||||
device int32_t * dst,
|
||||
threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
// bitonic sort
|
||||
const int col = tpitg[0];
|
||||
const int ib = tgpig[0] / args.ne01;
|
||||
|
||||
const int i00 = ib*ntg.x;
|
||||
const int i01 = tgpig[0] % args.ne01;
|
||||
const int i02 = tgpig[1];
|
||||
const int i03 = tgpig[2];
|
||||
|
||||
device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03);
|
||||
|
||||
// initialize indices
|
||||
shmem_i32[col] = i00 + col;
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (int k = 2; k <= ntg.x; k *= 2) {
|
||||
for (int j = k / 2; j > 0; j /= 2) {
|
||||
int ixj = col ^ j;
|
||||
if (ixj > col) {
|
||||
if ((col & k) == 0) {
|
||||
if (shmem_i32[col] >= args.ne00 ||
|
||||
(shmem_i32[ixj] < args.ne00 && (order == GGML_SORT_ORDER_ASC ?
|
||||
src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]] :
|
||||
src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]]))
|
||||
) {
|
||||
SWAP(shmem_i32[col], shmem_i32[ixj]);
|
||||
}
|
||||
} else {
|
||||
if (shmem_i32[ixj] >= args.ne00 ||
|
||||
(shmem_i32[col] < args.ne00 && (order == GGML_SORT_ORDER_ASC ?
|
||||
src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]] :
|
||||
src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]]))
|
||||
) {
|
||||
SWAP(shmem_i32[col], shmem_i32[ixj]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
}
|
||||
|
||||
const int64_t i0 = ib*args.top_k;
|
||||
|
||||
// copy the result to dst without the padding
|
||||
if (i0 + col < args.ne0 && col < args.top_k) {
|
||||
dst += i0 + args.ne0*i01 + args.ne0*args.ne1*i02 + args.ne0*args.ne1*args.ne2*i03;
|
||||
|
||||
dst[col] = shmem_i32[col];
|
||||
}
|
||||
}
|
||||
|
||||
template [[host_name("kernel_argsort_f32_i32_asc")]] kernel argsort_t kernel_argsort_f32_i32<GGML_SORT_ORDER_ASC>;
|
||||
template [[host_name("kernel_argsort_f32_i32_desc")]] kernel argsort_t kernel_argsort_f32_i32<GGML_SORT_ORDER_DESC>;
|
||||
|
||||
typedef void (argsort_merge_t)(
|
||||
constant ggml_metal_kargs_argsort_merge & args,
|
||||
device const char * src0,
|
||||
device const int32_t * tmp,
|
||||
device int32_t * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
template<ggml_sort_order order>
|
||||
kernel void kernel_argsort_merge_f32_i32(
|
||||
constant ggml_metal_kargs_argsort_merge & args,
|
||||
device const char * src0,
|
||||
device const int32_t * tmp,
|
||||
device int32_t * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int im = tgpig[0] / args.ne01;
|
||||
const int i01 = tgpig[0] % args.ne01;
|
||||
const int i02 = tgpig[1];
|
||||
const int i03 = tgpig[2];
|
||||
|
||||
const int start = im * (2 * args.len);
|
||||
|
||||
const int len0 = MIN(args.len, MAX(0, args.ne0 - (int)(start)));
|
||||
const int len1 = MIN(args.len, MAX(0, args.ne0 - (int)(start + args.len)));
|
||||
|
||||
const int total = len0 + len1;
|
||||
|
||||
device const int32_t * tmp0 = tmp + start
|
||||
+ i01*args.ne0
|
||||
+ i02*args.ne0*args.ne01
|
||||
+ i03*args.ne0*args.ne01*args.ne02;
|
||||
|
||||
device const int32_t * tmp1 = tmp0 + args.len;
|
||||
|
||||
dst += start
|
||||
+ i01*args.top_k
|
||||
+ i02*args.top_k*args.ne01
|
||||
+ i03*args.top_k*args.ne01*args.ne02;
|
||||
|
||||
device const float * src0_row = (device const float *)(src0
|
||||
+ args.nb01*i01
|
||||
+ args.nb02*i02
|
||||
+ args.nb03*i03);
|
||||
|
||||
if (total == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int chunk = (total + ntg.x - 1) / ntg.x;
|
||||
|
||||
const int k0 = tpitg.x * chunk;
|
||||
const int k1 = MIN(MIN(k0 + chunk, total), args.top_k);
|
||||
|
||||
if (k0 >= args.top_k) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (k0 >= total) {
|
||||
return;
|
||||
}
|
||||
|
||||
int low = k0 > len1 ? k0 - len1 : 0;
|
||||
int high = MIN(k0, len0);
|
||||
|
||||
// binary-search partition (i, j) such that i + j = k
|
||||
while (low < high) {
|
||||
const int mid = (low + high) >> 1;
|
||||
|
||||
const int32_t idx0 = tmp0[mid];
|
||||
const int32_t idx1 = tmp1[k0 - mid - 1];
|
||||
|
||||
const float val0 = src0_row[idx0];
|
||||
const float val1 = src0_row[idx1];
|
||||
|
||||
bool take_left;
|
||||
if (order == GGML_SORT_ORDER_ASC) {
|
||||
take_left = (val0 <= val1);
|
||||
} else {
|
||||
take_left = (val0 >= val1);
|
||||
}
|
||||
|
||||
if (take_left) {
|
||||
low = mid + 1;
|
||||
} else {
|
||||
high = mid;
|
||||
}
|
||||
}
|
||||
|
||||
int i = low;
|
||||
int j = k0 - i;
|
||||
|
||||
// keep the merge fronts into registers
|
||||
int32_t idx0 = 0;
|
||||
float val0 = 0.0f;
|
||||
if (i < len0) {
|
||||
idx0 = tmp0[i];
|
||||
val0 = src0_row[idx0];
|
||||
}
|
||||
|
||||
int32_t idx1 = 0;
|
||||
float val1 = 0.0f;
|
||||
if (j < len1) {
|
||||
idx1 = tmp1[j];
|
||||
val1 = src0_row[idx1];
|
||||
}
|
||||
|
||||
for (int k = k0; k < k1; ++k) {
|
||||
int32_t out_idx;
|
||||
|
||||
if (i >= len0) {
|
||||
while (k < k1) {
|
||||
dst[k++] = tmp1[j++];
|
||||
}
|
||||
break;
|
||||
} else if (j >= len1) {
|
||||
while (k < k1) {
|
||||
dst[k++] = tmp0[i++];
|
||||
}
|
||||
break;
|
||||
} else {
|
||||
bool take_left;
|
||||
|
||||
if (order == GGML_SORT_ORDER_ASC) {
|
||||
take_left = (val0 <= val1);
|
||||
} else {
|
||||
take_left = (val0 >= val1);
|
||||
}
|
||||
|
||||
if (take_left) {
|
||||
out_idx = idx0;
|
||||
++i;
|
||||
if (i < len0) {
|
||||
idx0 = tmp0[i];
|
||||
val0 = src0_row[idx0];
|
||||
}
|
||||
} else {
|
||||
out_idx = idx1;
|
||||
++j;
|
||||
if (j < len1) {
|
||||
idx1 = tmp1[j];
|
||||
val1 = src0_row[idx1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
dst[k] = out_idx;
|
||||
}
|
||||
}
|
||||
|
||||
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>;
|
||||
@@ -1,228 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
// OP: 0 - add, 1 - sub, 2 - mul, 3 - div
|
||||
constant short FC_bin_op [[function_constant(FC_BIN + 0)]];
|
||||
constant short FC_bin_f [[function_constant(FC_BIN + 1)]];
|
||||
constant bool FC_bin_rb [[function_constant(FC_BIN + 2)]];
|
||||
constant bool FC_bin_cb [[function_constant(FC_BIN + 3)]];
|
||||
|
||||
template <typename T0, typename T1, typename T>
|
||||
kernel void kernel_bin_fuse_impl(
|
||||
constant ggml_metal_kargs_bin & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
#define FC_OP FC_bin_op
|
||||
#define FC_F FC_bin_f
|
||||
#define FC_RB FC_bin_rb
|
||||
#define FC_CB FC_bin_cb
|
||||
|
||||
if (FC_RB) {
|
||||
// row broadcast
|
||||
const uint i0 = tgpig.y*args.ne00 + tgpig.x;
|
||||
const uint i1 = FC_CB ? tgpig.x%args.ne10 : tgpig.x;
|
||||
|
||||
device const T0 * src0_row = (device const T0 *) (src0);
|
||||
device T * dst_row = (device T *) (dst);
|
||||
|
||||
if (FC_F == 1) {
|
||||
device const T1 * src1_row = (device const T1 *) (src1 + args.o1[0]);
|
||||
|
||||
if (FC_OP == 0) {
|
||||
dst_row[i0] = src0_row[i0] + src1_row[i1];
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
dst_row[i0] = src0_row[i0] - src1_row[i1];
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
dst_row[i0] = src0_row[i0] * src1_row[i1];
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
dst_row[i0] = src0_row[i0] / src1_row[i1];
|
||||
}
|
||||
} else {
|
||||
T0 res = src0_row[i0];
|
||||
|
||||
if (FC_OP == 0) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res += ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res -= ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res *= ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res /= ((device const T1 *) (src1 + args.o1[j]))[i1];
|
||||
}
|
||||
}
|
||||
|
||||
dst_row[i0] = res;
|
||||
}
|
||||
} else {
|
||||
const int i03 = tgpig.z;
|
||||
const int i02 = tgpig.y;
|
||||
const int i01 = tgpig.x;
|
||||
|
||||
if (i01 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int i13 = i03%args.ne13;
|
||||
const int i12 = i02%args.ne12;
|
||||
const int i11 = i01%args.ne11;
|
||||
|
||||
device const T0 * src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + args.offs);
|
||||
device T * dst_ptr = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1 + args.offs);
|
||||
|
||||
if (FC_F == 1) {
|
||||
device const T1 * src1_ptr = (device const T1 *) (src1 + args.o1[0] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const int i10 = FC_CB ? i0%args.ne10 : i0;
|
||||
|
||||
if (FC_OP == 0) {
|
||||
dst_ptr[i0] = src0_ptr[i0] + src1_ptr[i10];
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
dst_ptr[i0] = src0_ptr[i0] - src1_ptr[i10];
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
dst_ptr[i0] = src0_ptr[i0] * src1_ptr[i10];
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
dst_ptr[i0] = src0_ptr[i0] / src1_ptr[i10];
|
||||
}
|
||||
}
|
||||
} else {
|
||||
device const T1 * src1_ptr[8];
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
src1_ptr[j] = (device const T1 *) (src1 + args.o1[j] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
|
||||
}
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const int i10 = FC_CB ? i0%args.ne10 : i0;
|
||||
|
||||
T res = src0_ptr[i0];
|
||||
|
||||
if (FC_OP == 0) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res += src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 1) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res -= src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 2) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res *= src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_OP == 3) {
|
||||
FOR_UNROLL (short j = 0; j < FC_F; ++j) {
|
||||
res /= src1_ptr[j][i10];
|
||||
}
|
||||
}
|
||||
|
||||
dst_ptr[i0] = res;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#undef FC_OP
|
||||
#undef FC_F
|
||||
#undef FC_RB
|
||||
#undef FC_CB
|
||||
}
|
||||
|
||||
typedef decltype(kernel_bin_fuse_impl<float, float, float>) kernel_bin_fuse_t;
|
||||
|
||||
template [[host_name("kernel_bin_fuse_f32_f32_f32")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float, float, float>;
|
||||
template [[host_name("kernel_bin_fuse_f32_f32_f32_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float4, float4, float4>;
|
||||
template [[host_name("kernel_bin_fuse_f16_f16_f16")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<half, half, half>;
|
||||
template [[host_name("kernel_bin_fuse_f16_f16_f16_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<half4, half4, half4>;
|
||||
|
||||
kernel void kernel_add_id(
|
||||
constant ggml_metal_kargs_add_id & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i1 = tgpig.x;
|
||||
const int i2 = tgpig.y;
|
||||
|
||||
const int i11 = *((device const int32_t *) (src2 + i1*sizeof(int32_t) + i2*args.nb21));
|
||||
|
||||
const size_t nb1 = args.ne0 * sizeof(float);
|
||||
const size_t nb2 = args.ne1 * nb1;
|
||||
|
||||
device float * dst_row = (device float *)((device char *)dst + i1*nb1 + i2*nb2);
|
||||
device const float * src0_row = (device const float *)((device char *)src0 + i1*args.nb01 + i2*args.nb02);
|
||||
device const float * src1_row = (device const float *)((device char *)src1 + i11*args.nb11);
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
dst_row[i0] = src0_row[i0] + src1_row[i0];
|
||||
}
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_repeat(
|
||||
constant ggml_metal_kargs_repeat & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = tgpig.x;
|
||||
|
||||
const int i03 = i3%args.ne03;
|
||||
const int i02 = i2%args.ne02;
|
||||
const int i01 = i1%args.ne01;
|
||||
|
||||
device const char * src0_ptr = src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01;
|
||||
device char * dst_ptr = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1;
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const int i00 = i0%args.ne00;
|
||||
*((device T *)(dst_ptr + i0*args.nb0)) = *((device T *)(src0_ptr + i00*args.nb00));
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_repeat<float>) kernel_repeat_t;
|
||||
|
||||
template [[host_name("kernel_repeat_f32")]] kernel kernel_repeat_t kernel_repeat<float>;
|
||||
template [[host_name("kernel_repeat_f16")]] kernel kernel_repeat_t kernel_repeat<half>;
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_repeat_bf16")]] kernel kernel_repeat_t kernel_repeat<bfloat>;
|
||||
#endif
|
||||
template [[host_name("kernel_repeat_i32")]] kernel kernel_repeat_t kernel_repeat<int>;
|
||||
template [[host_name("kernel_repeat_i16")]] kernel kernel_repeat_t kernel_repeat<short>;
|
||||
@@ -1,126 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
#include "ggml-metal-impl.h"
|
||||
|
||||
#include <metal_stdlib>
|
||||
|
||||
#ifdef GGML_METAL_HAS_TENSOR
|
||||
#include <metal_tensor>
|
||||
|
||||
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
|
||||
#endif
|
||||
|
||||
using namespace metal;
|
||||
|
||||
#define MAX(x, y) ((x) > (y) ? (x) : (y))
|
||||
#define MIN(x, y) ((x) < (y) ? (x) : (y))
|
||||
#define SWAP(x, y) { auto tmp = (x); (x) = (y); (y) = tmp; }
|
||||
|
||||
#define PAD2(x, n) (((x) + (n) - 1) & ~((n) - 1))
|
||||
|
||||
#define FOR_UNROLL(x) _Pragma("clang loop unroll(full)") for (x)
|
||||
|
||||
#define N_SIMDWIDTH 32 // assuming SIMD group size is 32
|
||||
|
||||
// ref: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
|
||||
//
|
||||
// cmd:
|
||||
// .../usr/bin/metal -dM -E -c ggml/src/ggml-metal/kernels/<src>.metal
|
||||
// .../usr/bin/metal -dM -E -c -target air64-apple-ios14.0 ggml/src/ggml-metal/kernels/<src>.metal
|
||||
//
|
||||
#if __METAL_VERSION__ < 310 && defined(GGML_METAL_HAS_BF16)
|
||||
#undef GGML_METAL_HAS_BF16
|
||||
#endif
|
||||
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
typedef matrix<bfloat, 4, 4> bfloat4x4;
|
||||
typedef matrix<bfloat, 2, 4> bfloat2x4;
|
||||
#endif
|
||||
|
||||
constexpr constant static float kvalues_iq4nl_f[16] = {
|
||||
-127.f, -104.f, -83.f, -65.f, -49.f, -35.f, -22.f, -10.f, 1.f, 13.f, 25.f, 38.f, 53.f, 69.f, 89.f, 113.f
|
||||
};
|
||||
|
||||
constexpr constant static float kvalues_mxfp4_f[16] = {
|
||||
0, .5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f, -0, -.5f, -1.f, -1.5f, -2.f, -3.f, -4.f, -6.f
|
||||
};
|
||||
|
||||
static inline int best_index_int8(int n, constant float * val, float x) {
|
||||
if (x <= val[0]) return 0;
|
||||
if (x >= val[n-1]) return n-1;
|
||||
int ml = 0, mu = n-1;
|
||||
while (mu-ml > 1) {
|
||||
int mav = (ml+mu)/2;
|
||||
if (x < val[mav]) mu = mav; else ml = mav;
|
||||
}
|
||||
return x - val[mu-1] < val[mu] - x ? mu-1 : mu;
|
||||
}
|
||||
|
||||
static inline float e8m0_to_fp32(uint8_t x) {
|
||||
uint32_t bits;
|
||||
|
||||
if (x == 0) {
|
||||
bits = 0x00400000;
|
||||
} else {
|
||||
bits = (uint32_t) x << 23;
|
||||
}
|
||||
|
||||
return as_type<float>(bits);
|
||||
}
|
||||
|
||||
static inline float dot(float x, float y) {
|
||||
return x*y;
|
||||
}
|
||||
|
||||
static inline float sum(float x) {
|
||||
return x;
|
||||
}
|
||||
|
||||
static inline float sum(float4 x) {
|
||||
return x[0] + x[1] + x[2] + x[3];
|
||||
}
|
||||
|
||||
enum ggml_sort_order {
|
||||
GGML_SORT_ORDER_ASC,
|
||||
GGML_SORT_ORDER_DESC,
|
||||
};
|
||||
|
||||
constant float GELU_COEF_A = 0.044715f;
|
||||
constant float GELU_QUICK_COEF = -1.702f;
|
||||
constant float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
|
||||
constant float SQRT_2_INV = 0.70710678118654752440084436210484f;
|
||||
|
||||
// based on Abramowitz and Stegun formula 7.1.26 or similar Hastings' approximation
|
||||
// ref: https://www.johndcook.com/blog/python_erf/
|
||||
constant float p_erf = 0.3275911f;
|
||||
constant float a1_erf = 0.254829592f;
|
||||
constant float a2_erf = -0.284496736f;
|
||||
constant float a3_erf = 1.421413741f;
|
||||
constant float a4_erf = -1.453152027f;
|
||||
constant float a5_erf = 1.061405429f;
|
||||
|
||||
template<typename T>
|
||||
inline T erf_approx(T x) {
|
||||
T sign_x = sign(x);
|
||||
x = fabs(x);
|
||||
T t = 1.0f / (1.0f + p_erf * x);
|
||||
T y = 1.0f - (((((a5_erf * t + a4_erf) * t) + a3_erf) * t + a2_erf) * t + a1_erf) * t * exp(-x * x);
|
||||
return sign_x * y;
|
||||
}
|
||||
|
||||
template<typename T> T elu_approx(T x);
|
||||
|
||||
template<> inline float elu_approx<float>(float x) {
|
||||
return (x > 0.f) ? x : (exp(x) - 1);
|
||||
}
|
||||
|
||||
template<> inline float4 elu_approx<float4>(float4 x) {
|
||||
float4 res;
|
||||
|
||||
res[0] = (x[0] > 0.0f) ? x[0] : (exp(x[0]) - 1.0f);
|
||||
res[1] = (x[1] > 0.0f) ? x[1] : (exp(x[1]) - 1.0f);
|
||||
res[2] = (x[2] > 0.0f) ? x[2] : (exp(x[2]) - 1.0f);
|
||||
res[3] = (x[3] > 0.0f) ? x[3] : (exp(x[3]) - 1.0f);
|
||||
|
||||
return res;
|
||||
}
|
||||
@@ -1,723 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
typedef void (im2col_t)(
|
||||
constant ggml_metal_kargs_im2col & args,
|
||||
device const float * x,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
template <typename T>
|
||||
kernel void kernel_im2col(
|
||||
constant ggml_metal_kargs_im2col & args,
|
||||
device const float * x,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
// const int64_t IC = tgpg[0];
|
||||
const int64_t OH = tgpg[1];
|
||||
const int64_t OW = tgpg[2];
|
||||
|
||||
const int64_t KH = ntg[1];
|
||||
const int64_t KW = ntg[2];
|
||||
|
||||
int64_t in = tpitg[0];
|
||||
const int64_t ikh = tpitg[1];
|
||||
const int64_t ikw = tpitg[2];
|
||||
|
||||
const int64_t iic = tgpig[0];
|
||||
const int64_t ioh = tgpig[1];
|
||||
const int64_t iow = tgpig[2];
|
||||
|
||||
const int64_t iiw = iow*args.s0 + ikw*args.d0 - args.p0;
|
||||
const int64_t iih = ioh*args.s1 + ikh*args.d1 - args.p1;
|
||||
|
||||
int64_t offset_dst = (in*OH*OW + ioh*OW + iow)*args.CHW + (iic*(KH*KW) + ikh*KW + ikw);
|
||||
|
||||
device T * pdst = (device T *) (dst);
|
||||
|
||||
if (iih < 0 || iih >= args.IH || iiw < 0 || iiw >= args.IW) {
|
||||
while (in < args.N) {
|
||||
pdst[offset_dst] = 0.0f;
|
||||
offset_dst += ntg[0]*args.CHW*OH*OW;
|
||||
|
||||
in += ntg[0];
|
||||
}
|
||||
} else {
|
||||
int64_t offset_src = in*args.ofs0 + iic*args.ofs1 + iih*args.IW + iiw;
|
||||
|
||||
while (in < args.N) {
|
||||
pdst[offset_dst] = x[offset_src];
|
||||
|
||||
offset_dst += ntg[0]*args.CHW*OH*OW;
|
||||
offset_src += ntg[0]*args.ofs0;
|
||||
|
||||
in += ntg[0];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template [[host_name("kernel_im2col_f32")]] kernel im2col_t kernel_im2col<float>;
|
||||
template [[host_name("kernel_im2col_f16")]] kernel im2col_t kernel_im2col<half>;
|
||||
|
||||
// TODO: optimize
|
||||
typedef void (im2col_ext_t)(
|
||||
constant ggml_metal_kargs_im2col & args,
|
||||
device const float * x,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
template <typename T>
|
||||
kernel void kernel_im2col_ext(
|
||||
constant ggml_metal_kargs_im2col & args,
|
||||
device const float * x,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]], // tgpg[0] = D x IC x KH x KW, CHW = IC x KH x KW
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) { // [M, 1, 1]
|
||||
const int64_t KHW = (int64_t)args.KHW;
|
||||
|
||||
const int64_t d = tgpig[0] / args.CHW;
|
||||
const int64_t chw = tgpig[0] % args.CHW;
|
||||
const int64_t tgpig_0 = chw / KHW; // 0 ~ (IC - 1)
|
||||
const int64_t HW = tgpig[0] % KHW;
|
||||
|
||||
const int64_t tpitg_0 = (d * ntg[0]) + tpitg[0];
|
||||
if (tpitg_0 >= args.N) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t tpitg_1 = HW / args.KW;
|
||||
const int64_t tpitg_2 = HW % args.KW;
|
||||
|
||||
const int64_t iiw = tgpig[2] * args.s0 + tpitg_2 * args.d0 - args.p0;
|
||||
const int64_t iih = tgpig[1] * args.s1 + tpitg_1 * args.d1 - args.p1;
|
||||
|
||||
const int64_t offset_dst =
|
||||
(tpitg_0 * tgpg[1] * tgpg[2] + tgpig[1] * tgpg[2] + tgpig[2]) * args.CHW +
|
||||
(tgpig_0 * KHW + tpitg_1 * args.KW + tpitg_2);
|
||||
|
||||
device T * pdst = (device T *) (dst);
|
||||
|
||||
if (iih < 0 || iih >= args.IH || iiw < 0 || iiw >= args.IW) {
|
||||
pdst[offset_dst] = 0.0f;
|
||||
} else {
|
||||
const int64_t offset_src = tpitg_0 * args.ofs0 + tgpig_0 * args.ofs1;
|
||||
pdst[offset_dst] = x[offset_src + iih * args.IW + iiw];
|
||||
}
|
||||
}
|
||||
|
||||
template [[host_name("kernel_im2col_ext_f32")]] kernel im2col_ext_t kernel_im2col_ext<float>;
|
||||
template [[host_name("kernel_im2col_ext_f16")]] kernel im2col_ext_t kernel_im2col_ext<half>;
|
||||
|
||||
template <typename T>
|
||||
kernel void kernel_col2im_1d(
|
||||
constant ggml_metal_kargs_col2im_1d & args,
|
||||
device const T * col,
|
||||
device T * dst,
|
||||
uint tgpig [[threadgroup_position_in_grid]],
|
||||
uint tpitg [[thread_position_in_threadgroup]],
|
||||
uint ntg [[threads_per_threadgroup]]) {
|
||||
|
||||
const int idx = tgpig * ntg + tpitg;
|
||||
if (idx >= args.T_out * args.OC) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int t_out = idx % args.T_out;
|
||||
const int oc = idx / args.T_out;
|
||||
const int t_abs = t_out + args.p0; // absolute position in uncropped signal
|
||||
|
||||
int t_in_min = (t_abs - args.K + args.s0) / args.s0; // ceil((t_abs - K + 1) / s0)
|
||||
if (t_in_min < 0) {
|
||||
t_in_min = 0;
|
||||
}
|
||||
int t_in_max = t_abs / args.s0;
|
||||
if (t_in_max >= args.T_in) {
|
||||
t_in_max = args.T_in - 1;
|
||||
}
|
||||
|
||||
float sum = 0.0f;
|
||||
for (int t_in = t_in_min; t_in <= t_in_max; t_in++) {
|
||||
const int k = t_abs - t_in * args.s0;
|
||||
sum += float(col[(oc * args.K + k) + t_in * args.K_OC]);
|
||||
}
|
||||
|
||||
dst[t_out + oc * args.T_out] = T(sum);
|
||||
}
|
||||
|
||||
template [[host_name("kernel_col2im_1d_f32")]] kernel void kernel_col2im_1d<float>(constant ggml_metal_kargs_col2im_1d &, device const float *, device float *, uint, uint, uint);
|
||||
template [[host_name("kernel_col2im_1d_f16")]] kernel void kernel_col2im_1d<half>(constant ggml_metal_kargs_col2im_1d &, device const half *, device half *, uint, uint, uint);
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_col2im_1d_bf16")]] kernel void kernel_col2im_1d<bfloat>(constant ggml_metal_kargs_col2im_1d &, device const bfloat *, device bfloat *, uint, uint, uint);
|
||||
#endif
|
||||
|
||||
template <typename TK>
|
||||
kernel void kernel_conv_2d(
|
||||
constant ggml_metal_kargs_conv_2d & args,
|
||||
device const char * weights,
|
||||
device const char * src,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const uint threads_per_tg = ntg.x * ntg.y * ntg.z;
|
||||
const uint tg_index = (tgpig.z * tgpg.y + tgpig.y) * tgpg.x + tgpig.x;
|
||||
const uint local_thread = tpitg.z * (ntg.x * ntg.y) + tpitg.y * ntg.x + tpitg.x;
|
||||
const uint thread_index = tg_index * threads_per_tg + local_thread;
|
||||
const uint64_t total_threads = (uint64_t) threads_per_tg * tgpg.x * tgpg.y * tgpg.z;
|
||||
const uint64_t total_outputs = (uint64_t) args.N * args.OC * args.OH * args.OW;
|
||||
|
||||
for (uint64_t index = thread_index; index < total_outputs; index += total_threads) {
|
||||
uint64_t tmp = index;
|
||||
|
||||
const int32_t ow = tmp % args.OW; tmp /= args.OW;
|
||||
const int32_t oh = tmp % args.OH; tmp /= args.OH;
|
||||
const int32_t oc = tmp % args.OC; tmp /= args.OC;
|
||||
const int32_t n = tmp;
|
||||
|
||||
float acc = 0.0f;
|
||||
|
||||
const int32_t base_x = ow*args.s0 - args.p0;
|
||||
const int32_t base_y = oh*args.s1 - args.p1;
|
||||
|
||||
int32_t ky_start = 0;
|
||||
if (base_y < 0) {
|
||||
ky_start = (-base_y + args.d1 - 1)/args.d1;
|
||||
}
|
||||
int32_t ky_end = args.KH;
|
||||
const int32_t y_max = args.IH - 1 - base_y;
|
||||
if (y_max < 0) {
|
||||
ky_end = ky_start;
|
||||
} else if (base_y + (args.KH - 1)*args.d1 >= args.IH) {
|
||||
ky_end = min(ky_end, y_max/args.d1 + 1);
|
||||
}
|
||||
|
||||
int32_t kx_start = 0;
|
||||
if (base_x < 0) {
|
||||
kx_start = (-base_x + args.d0 - 1)/args.d0;
|
||||
}
|
||||
int32_t kx_end = args.KW;
|
||||
const int32_t x_max = args.IW - 1 - base_x;
|
||||
if (x_max < 0) {
|
||||
kx_end = kx_start;
|
||||
} else if (base_x + (args.KW - 1)*args.d0 >= args.IW) {
|
||||
kx_end = min(kx_end, x_max/args.d0 + 1);
|
||||
}
|
||||
|
||||
if (ky_start < ky_end && kx_start < kx_end) {
|
||||
const uint64_t src_base_n = (uint64_t) n * args.nb13;
|
||||
const uint64_t w_base_oc = (uint64_t) oc * args.nb03;
|
||||
|
||||
for (int32_t ic = 0; ic < args.IC; ++ic) {
|
||||
const uint64_t src_base_nc = src_base_n + (uint64_t) ic * args.nb12;
|
||||
const uint64_t w_base_ocic = w_base_oc + (uint64_t) ic * args.nb02;
|
||||
|
||||
for (int32_t ky = ky_start; ky < ky_end; ++ky) {
|
||||
const int32_t iy = base_y + ky*args.d1;
|
||||
const uint64_t src_base_row = src_base_nc + (uint64_t) iy * args.nb11;
|
||||
const uint64_t w_base_row = w_base_ocic + (uint64_t) ky * args.nb01;
|
||||
|
||||
for (int32_t kx = kx_start; kx < kx_end; ++kx) {
|
||||
const int32_t ix = base_x + kx*args.d0;
|
||||
const uint64_t src_offs = src_base_row + (uint64_t) ix * args.nb10;
|
||||
const uint64_t w_offs = w_base_row + (uint64_t) kx * args.nb00;
|
||||
|
||||
const float x = *(device const float *)(src + src_offs);
|
||||
const float w = (float) (*(device const TK *)(weights + w_offs));
|
||||
|
||||
acc += x * w;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const uint64_t dst_offs =
|
||||
(uint64_t) n * args.nb3 +
|
||||
(uint64_t) oc * args.nb2 +
|
||||
(uint64_t) oh * args.nb1 +
|
||||
(uint64_t) ow * args.nb0;
|
||||
|
||||
*(device float *)(dst + dst_offs) = acc;
|
||||
}
|
||||
}
|
||||
|
||||
template [[host_name("kernel_conv_2d_f32_f32")]]
|
||||
kernel void kernel_conv_2d<float>(
|
||||
constant ggml_metal_kargs_conv_2d & args,
|
||||
device const char * weights,
|
||||
device const char * src,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
template [[host_name("kernel_conv_2d_f16_f32")]]
|
||||
kernel void kernel_conv_2d<half>(
|
||||
constant ggml_metal_kargs_conv_2d & args,
|
||||
device const char * weights,
|
||||
device const char * src,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
typedef void (conv_transpose_1d_t)(
|
||||
constant ggml_metal_kargs_conv_transpose_1d & args,
|
||||
device const float * src0,
|
||||
device const float * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]]);
|
||||
|
||||
template <typename T>
|
||||
kernel void kernel_conv_transpose_1d(
|
||||
constant ggml_metal_kargs_conv_transpose_1d & args,
|
||||
device const T * src0,
|
||||
device const float * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]]) {
|
||||
|
||||
// For output position j on the time axis, only input positions
|
||||
// i such that i*s0 <= j < i*s0 + K
|
||||
// contribute -- i.e. i in [ceil((j - K + 1)/s0), floor(j/s0)]
|
||||
// intersected with [0, IL-1]. That's at most ceil(K/s0) values
|
||||
// (typically 2 for stride==K/2 transposed convs).
|
||||
const int32_t j = tgpig[0];
|
||||
const int32_t s0 = args.s0;
|
||||
const int32_t K = args.K;
|
||||
const int32_t IL = args.IL;
|
||||
|
||||
int32_t i_min;
|
||||
{
|
||||
int32_t a = j - K + 1;
|
||||
i_min = a <= 0 ? 0 : (a + s0 - 1) / s0; // ceil(a/s0) for a>0
|
||||
}
|
||||
int32_t i_max = j / s0;
|
||||
if (i_max > IL - 1) i_max = IL - 1;
|
||||
|
||||
float v = 0.0f;
|
||||
if (i_min <= i_max) {
|
||||
for (int64_t c = 0; c < args.IC; c++) {
|
||||
const int32_t kernel_offset = c * tgpg[1] * K + K * tgpig[1];
|
||||
const int32_t input_offset = c * IL;
|
||||
|
||||
for (int32_t i = i_min; i <= i_max; i++) {
|
||||
v += float(src0[kernel_offset + j - i * s0]) * src1[input_offset + i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
device float * dst_ptr = (device float *) (dst + tgpig[0] * args.nb0 + tgpig[1] * args.nb1);
|
||||
|
||||
dst_ptr[0] = v;
|
||||
}
|
||||
|
||||
template [[host_name("kernel_conv_transpose_1d_f32_f32")]]
|
||||
kernel void kernel_conv_transpose_1d<float>(
|
||||
constant ggml_metal_kargs_conv_transpose_1d & args,
|
||||
device const float * src0,
|
||||
device const float * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]]);
|
||||
|
||||
template [[host_name("kernel_conv_transpose_1d_f16_f32")]]
|
||||
kernel void kernel_conv_transpose_1d<half>(
|
||||
constant ggml_metal_kargs_conv_transpose_1d & args,
|
||||
device const half * src0,
|
||||
device const float * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]]);
|
||||
|
||||
|
||||
typedef void (conv_transpose_2d_t)(
|
||||
constant ggml_metal_kargs_conv_transpose_2d & args,
|
||||
device const float * src0,
|
||||
device const float * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]]);
|
||||
|
||||
template <typename T>
|
||||
kernel void kernel_conv_transpose_2d(
|
||||
constant ggml_metal_kargs_conv_transpose_2d & args,
|
||||
device const T * src0,
|
||||
device const float * src1,
|
||||
device char * dst,
|
||||
threadgroup float * shared_sum [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int64_t out_x = tgpig[0];
|
||||
const int64_t out_y = tgpig[1];
|
||||
const int64_t out_c = tgpig[2];
|
||||
|
||||
const int64_t kw = tpitg[0];
|
||||
const int64_t kh = tpitg[1];
|
||||
|
||||
float v = 0.0f;
|
||||
|
||||
for (int64_t in_c = 0; in_c < args.IC; in_c++) {
|
||||
int64_t in_y = out_y - kh;
|
||||
|
||||
if (in_y < 0 || in_y % args.s0) continue;
|
||||
|
||||
in_y /= args.s0;
|
||||
|
||||
if (in_y >= args.IH) continue;
|
||||
|
||||
int64_t in_x = out_x - kw;
|
||||
|
||||
if (in_x < 0 || in_x % args.s0) continue;
|
||||
|
||||
in_x /= args.s0;
|
||||
|
||||
if (in_x >= args.IW) continue;
|
||||
|
||||
const int64_t input_idx = (args.IW * args.IH) * in_c + (args.IW) * in_y + in_x;
|
||||
const int64_t kernel_idx = (args.KH * args.KW * args.OC) * in_c + (args.KH * args.KW) * out_c + (args.KW) * kh + kw;
|
||||
|
||||
v += (float)src0[kernel_idx] * src1[input_idx];
|
||||
}
|
||||
|
||||
const uint tid = tpitg.y * ntg.x + tpitg.x;
|
||||
shared_sum[tid] = v;
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tid == 0) {
|
||||
float total = 0.0f;
|
||||
const uint num_threads = ntg.x * ntg.y;
|
||||
for (uint i = 0; i < num_threads; i++) {
|
||||
total += shared_sum[i];
|
||||
}
|
||||
|
||||
device float * dst_ptr = (device float *) (dst + out_x*args.nb0 + out_y * args.nb1 + out_c*args.nb2);
|
||||
dst_ptr[0] = total;
|
||||
}
|
||||
}
|
||||
|
||||
template [[host_name("kernel_conv_transpose_2d_f32_f32")]]
|
||||
kernel void kernel_conv_transpose_2d<float>(
|
||||
constant ggml_metal_kargs_conv_transpose_2d & args,
|
||||
device const float * src0,
|
||||
device const float * src1,
|
||||
device char * dst,
|
||||
threadgroup float * shared_sum [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
template [[host_name("kernel_conv_transpose_2d_f16_f32")]]
|
||||
kernel void kernel_conv_transpose_2d<half>(
|
||||
constant ggml_metal_kargs_conv_transpose_2d & args,
|
||||
device const half * src0,
|
||||
device const float * src1,
|
||||
device char * dst,
|
||||
threadgroup float * shared_sum [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
// grid: x = C tile, y = OH, z = OW * N (for channel-contiguous layouts)
|
||||
template <typename TK>
|
||||
kernel void kernel_conv_2d_dw_tiled(
|
||||
constant ggml_metal_kargs_conv_2d_dw & args,
|
||||
device const char * weights,
|
||||
device const char * src,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int32_t c = (int32_t)(tgpig.x * ntg.x + tpitg.x);
|
||||
if (c >= args.C) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int32_t oh = tgpig.y;
|
||||
const int32_t own = tgpig.z;
|
||||
const int32_t ow = own % args.OW;
|
||||
const int32_t n = own / args.OW;
|
||||
|
||||
const int32_t base_y = oh*args.s1 - args.p1;
|
||||
|
||||
int32_t ky_start = 0;
|
||||
if (base_y < 0) {
|
||||
ky_start = (-base_y + args.d1 - 1)/args.d1;
|
||||
}
|
||||
int32_t ky_end = args.KH;
|
||||
const int32_t y_max = args.IH - 1 - base_y;
|
||||
if (y_max < 0) {
|
||||
ky_end = ky_start;
|
||||
} else if (base_y + (args.KH - 1)*args.d1 >= args.IH) {
|
||||
ky_end = min(ky_end, y_max/args.d1 + 1);
|
||||
}
|
||||
|
||||
const int32_t base_x = ow*args.s0 - args.p0;
|
||||
|
||||
int32_t kx_start = 0;
|
||||
if (base_x < 0) {
|
||||
kx_start = (-base_x + args.d0 - 1)/args.d0;
|
||||
}
|
||||
int32_t kx_end = args.KW;
|
||||
const int32_t x_max = args.IW - 1 - base_x;
|
||||
if (x_max < 0) {
|
||||
kx_end = kx_start;
|
||||
} else if (base_x + (args.KW - 1)*args.d0 >= args.IW) {
|
||||
kx_end = min(kx_end, x_max/args.d0 + 1);
|
||||
}
|
||||
|
||||
float acc = 0.0f;
|
||||
|
||||
if (ky_start < ky_end && kx_start < kx_end) {
|
||||
const uint64_t w_base = (uint64_t) c * args.nb02;
|
||||
const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12;
|
||||
|
||||
for (int32_t ky = ky_start; ky < ky_end; ++ky) {
|
||||
const int32_t iy = base_y + ky*args.d1;
|
||||
const uint64_t src_row = src_base + (uint64_t) iy * args.nb11;
|
||||
const uint64_t w_row = w_base + (uint64_t) ky * args.nb01;
|
||||
|
||||
for (int32_t kx = kx_start; kx < kx_end; ++kx) {
|
||||
const int32_t ix = base_x + kx*args.d0;
|
||||
const float x = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10);
|
||||
const float w = (float)(*(device const TK *)(weights + w_row + (uint64_t) kx * args.nb00));
|
||||
acc += x * w;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const uint64_t dst_offs =
|
||||
(uint64_t) n * args.nb3 +
|
||||
(uint64_t) c * args.nb2 +
|
||||
(uint64_t) oh * args.nb1 +
|
||||
(uint64_t) ow * args.nb0;
|
||||
|
||||
*(device float *)(dst + dst_offs) = acc;
|
||||
}
|
||||
|
||||
// grid: x = OW tile, y = OH, z = C * N (for spatially-contiguous layouts)
|
||||
template <typename TK>
|
||||
kernel void kernel_conv_2d_dw(
|
||||
constant ggml_metal_kargs_conv_2d_dw & args,
|
||||
device const char * weights,
|
||||
device const char * src,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int32_t oh = tgpig.y;
|
||||
const int32_t cn = tgpig.z;
|
||||
const int32_t c = cn % args.C;
|
||||
const int32_t n = cn / args.C;
|
||||
|
||||
const int32_t base_y = oh*args.s1 - args.p1;
|
||||
|
||||
int32_t ky_start = 0;
|
||||
if (base_y < 0) {
|
||||
ky_start = (-base_y + args.d1 - 1)/args.d1;
|
||||
}
|
||||
int32_t ky_end = args.KH;
|
||||
const int32_t y_max = args.IH - 1 - base_y;
|
||||
if (y_max < 0) {
|
||||
ky_end = ky_start;
|
||||
} else if (base_y + (args.KH - 1)*args.d1 >= args.IH) {
|
||||
ky_end = min(ky_end, y_max/args.d1 + 1);
|
||||
}
|
||||
|
||||
const uint64_t w_base = (uint64_t) c * args.nb02;
|
||||
const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12;
|
||||
|
||||
const int32_t ow = (int32_t)(tgpig.x * ntg.x + tpitg.x);
|
||||
if (ow >= args.OW) {
|
||||
return;
|
||||
}
|
||||
|
||||
float acc = 0.0f;
|
||||
|
||||
const int32_t base_x = ow*args.s0 - args.p0;
|
||||
|
||||
int32_t kx_start = 0;
|
||||
if (base_x < 0) {
|
||||
kx_start = (-base_x + args.d0 - 1)/args.d0;
|
||||
}
|
||||
int32_t kx_end = args.KW;
|
||||
const int32_t x_max = args.IW - 1 - base_x;
|
||||
if (x_max < 0) {
|
||||
kx_end = kx_start;
|
||||
} else if (base_x + (args.KW - 1)*args.d0 >= args.IW) {
|
||||
kx_end = min(kx_end, x_max/args.d0 + 1);
|
||||
}
|
||||
|
||||
if (ky_start < ky_end && kx_start < kx_end) {
|
||||
for (int32_t ky = ky_start; ky < ky_end; ++ky) {
|
||||
const int32_t iy = base_y + ky*args.d1;
|
||||
const uint64_t src_row = src_base + (uint64_t) iy * args.nb11;
|
||||
const uint64_t w_row = w_base + (uint64_t) ky * args.nb01;
|
||||
|
||||
for (int32_t kx = kx_start; kx < kx_end; ++kx) {
|
||||
const int32_t ix = base_x + kx*args.d0;
|
||||
const float x = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10);
|
||||
const float w = (float)(*(device const TK *)(weights + w_row + (uint64_t) kx * args.nb00));
|
||||
acc += x * w;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const uint64_t dst_offs =
|
||||
(uint64_t) n * args.nb3 +
|
||||
(uint64_t) c * args.nb2 +
|
||||
(uint64_t) oh * args.nb1 +
|
||||
(uint64_t) ow * args.nb0;
|
||||
|
||||
*(device float *)(dst + dst_offs) = acc;
|
||||
}
|
||||
|
||||
template [[host_name("kernel_conv_2d_dw_f32_f32")]]
|
||||
kernel void kernel_conv_2d_dw<float>(
|
||||
constant ggml_metal_kargs_conv_2d_dw & args,
|
||||
device const char * weights,
|
||||
device const char * src,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
template [[host_name("kernel_conv_2d_dw_f16_f32")]]
|
||||
kernel void kernel_conv_2d_dw<half>(
|
||||
constant ggml_metal_kargs_conv_2d_dw & args,
|
||||
device const char * weights,
|
||||
device const char * src,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
template [[host_name("kernel_conv_2d_dw_tiled_f32_f32")]]
|
||||
kernel void kernel_conv_2d_dw_tiled<float>(
|
||||
constant ggml_metal_kargs_conv_2d_dw & args,
|
||||
device const char * weights,
|
||||
device const char * src,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
template [[host_name("kernel_conv_2d_dw_tiled_f16_f32")]]
|
||||
kernel void kernel_conv_2d_dw_tiled<half>(
|
||||
constant ggml_metal_kargs_conv_2d_dw & args,
|
||||
device const char * weights,
|
||||
device const char * src,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]);
|
||||
|
||||
template <typename T>
|
||||
kernel void kernel_conv_3d(
|
||||
constant ggml_metal_kargs_conv_3d & args,
|
||||
device const char * src0, // Weights [IC * OC, KD, KH, KW]
|
||||
device const char * src1, // Inputs [IC * N, ID, IH, IW]
|
||||
device char * dst, // Outputs [OC * N, OD, OH, OW]
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]]) {
|
||||
|
||||
// 1. Un-flatten the spatial dimension from Grid X
|
||||
int64_t spatial_idx = tgpig.x * 32 + tpitg.x;
|
||||
|
||||
if (spatial_idx >= args.OW * args.OH * args.OD) {
|
||||
return; // Thread falls outside the spatial volume
|
||||
}
|
||||
|
||||
int64_t od = spatial_idx / (args.OW * args.OH);
|
||||
int64_t oh = (spatial_idx / args.OW) % args.OH;
|
||||
int64_t ow = spatial_idx % args.OW;
|
||||
|
||||
// 2. Map Y to Channels, Z to Batch
|
||||
int64_t oc = tgpig.y;
|
||||
int64_t batch_idx = tgpig.z;
|
||||
|
||||
// 3. Calculate anchor coordinates in the Input volume
|
||||
int64_t i_w_base = ow * args.s0 - args.p0;
|
||||
int64_t i_h_base = oh * args.s1 - args.p1;
|
||||
int64_t i_d_base = od * args.s2 - args.p2;
|
||||
|
||||
float sum = 0.0f;
|
||||
|
||||
// 4. Gather Loop (Iterate over Input Channels -> Depth -> Height -> Width)
|
||||
for (int64_t ic = 0; ic < args.IC; ++ic) {
|
||||
|
||||
// ggml packs batch and channel together in the 4th dimension
|
||||
int64_t src_cn_idx = batch_idx * args.IC + ic;
|
||||
int64_t w_cn_idx = oc * args.IC + ic;
|
||||
|
||||
for (int64_t kz = 0; kz < args.KD; ++kz) {
|
||||
int64_t id = i_d_base + kz * args.d2;
|
||||
if (id < 0 || id >= args.ID) continue; // Boundary check (Padding)
|
||||
|
||||
for (int64_t ky = 0; ky < args.KH; ++ky) {
|
||||
int64_t ih = i_h_base + ky * args.d1;
|
||||
if (ih < 0 || ih >= args.IH) continue;
|
||||
|
||||
for (int64_t kx = 0; kx < args.KW; ++kx) {
|
||||
int64_t iw = i_w_base + kx * args.d0;
|
||||
if (iw < 0 || iw >= args.IW) continue;
|
||||
|
||||
// Convert multi-dimensional coordinates to flat byte offsets
|
||||
int64_t w_idx = kx*args.nb00 + ky*args.nb01 + kz*args.nb02 + w_cn_idx*args.nb03;
|
||||
int64_t i_idx = iw*args.nb10 + ih*args.nb11 + id*args.nb12 + src_cn_idx*args.nb13;
|
||||
|
||||
// Dereference memory and cast weights to f32 if they were f16
|
||||
float w_val = (float)*(device const T*)((device const char*)src0 + w_idx);
|
||||
float i_val = *(device const float*)((device const char*)src1 + i_idx);
|
||||
|
||||
sum += w_val * i_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Write the accumulated value out to RAM
|
||||
int64_t dst_cn_idx = batch_idx * args.OC + oc;
|
||||
int64_t d_idx = ow*args.nb0 + oh*args.nb1 + od*args.nb2 + dst_cn_idx*args.nb3;
|
||||
|
||||
*(device float*)(dst + d_idx) = sum;
|
||||
}
|
||||
|
||||
// Explicit instantiations so the JIT compiler can find them by name
|
||||
template [[host_name("kernel_conv_3d_f32_f32")]]
|
||||
kernel void kernel_conv_3d<float>(
|
||||
constant ggml_metal_kargs_conv_3d & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]]);
|
||||
|
||||
// Explicit instantiation for f16 weights
|
||||
template [[host_name("kernel_conv_3d_f16_f32")]]
|
||||
kernel void kernel_conv_3d<half>(
|
||||
constant ggml_metal_kargs_conv_3d & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]]);
|
||||
@@ -1,719 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
#include "common.h"
|
||||
|
||||
#define GGML_COMMON_DECL_METAL
|
||||
#define GGML_COMMON_IMPL_METAL
|
||||
#if defined(GGML_METAL_EMBED_LIBRARY)
|
||||
__embed_ggml-common.h__
|
||||
#else
|
||||
#include "ggml-common.h"
|
||||
#endif
|
||||
|
||||
#define QK_NL 16 // shared by mul_mm and get_rows_q instantiations
|
||||
|
||||
// NOTE: this is not dequantizing - we are simply fitting the template
|
||||
template <typename type4x4>
|
||||
void dequantize_f32(device const float4x4 * src, short il, thread type4x4 & reg) {
|
||||
reg = (type4x4)(*src);
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_f32_t4(device const float4 * src, short il, thread type4 & reg) {
|
||||
reg = (type4)(*src);
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_f16(device const half4x4 * src, short il, thread type4x4 & reg) {
|
||||
reg = (type4x4)(*src);
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_f16_t4(device const half4 * src, short il, thread type4 & reg) {
|
||||
reg = (type4)(*(src));
|
||||
}
|
||||
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template <typename type4x4>
|
||||
void dequantize_bf16(device const bfloat4x4 * src, short il, thread type4x4 & reg) {
|
||||
reg = (type4x4)(*src);
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_bf16_t4(device const bfloat4 * src, short il, thread type4 & reg) {
|
||||
reg = (type4)(*(src));
|
||||
}
|
||||
#endif
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q1_0(device const block_q1_0 * xb, short il, thread type4x4 & reg) {
|
||||
device const uint8_t * qs = xb->qs;
|
||||
const float d = xb->d;
|
||||
const float neg_d = -d;
|
||||
|
||||
const int byte_offset = il * 2; // il*16 bits = il*2 bytes
|
||||
const uint8_t b0 = qs[byte_offset];
|
||||
const uint8_t b1 = qs[byte_offset + 1];
|
||||
|
||||
float4x4 reg_f;
|
||||
|
||||
reg_f[0][0] = select(neg_d, d, bool(b0 & 0x01));
|
||||
reg_f[0][1] = select(neg_d, d, bool(b0 & 0x02));
|
||||
reg_f[0][2] = select(neg_d, d, bool(b0 & 0x04));
|
||||
reg_f[0][3] = select(neg_d, d, bool(b0 & 0x08));
|
||||
reg_f[1][0] = select(neg_d, d, bool(b0 & 0x10));
|
||||
reg_f[1][1] = select(neg_d, d, bool(b0 & 0x20));
|
||||
reg_f[1][2] = select(neg_d, d, bool(b0 & 0x40));
|
||||
reg_f[1][3] = select(neg_d, d, bool(b0 & 0x80));
|
||||
|
||||
reg_f[2][0] = select(neg_d, d, bool(b1 & 0x01));
|
||||
reg_f[2][1] = select(neg_d, d, bool(b1 & 0x02));
|
||||
reg_f[2][2] = select(neg_d, d, bool(b1 & 0x04));
|
||||
reg_f[2][3] = select(neg_d, d, bool(b1 & 0x08));
|
||||
reg_f[3][0] = select(neg_d, d, bool(b1 & 0x10));
|
||||
reg_f[3][1] = select(neg_d, d, bool(b1 & 0x20));
|
||||
reg_f[3][2] = select(neg_d, d, bool(b1 & 0x40));
|
||||
reg_f[3][3] = select(neg_d, d, bool(b1 & 0x80));
|
||||
|
||||
reg = (type4x4) reg_f;
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_q1_0_t4(device const block_q1_0 * xb, short il, thread type4 & reg) {
|
||||
const float d = xb->d;
|
||||
const float neg_d = -d;
|
||||
const int base = il * 4;
|
||||
const uint8_t byte = xb->qs[base / 8];
|
||||
const int s = base % 8;
|
||||
|
||||
float4 reg_f;
|
||||
reg_f[0] = select(neg_d, d, bool((byte >> (s )) & 1));
|
||||
reg_f[1] = select(neg_d, d, bool((byte >> (s + 1)) & 1));
|
||||
reg_f[2] = select(neg_d, d, bool((byte >> (s + 2)) & 1));
|
||||
reg_f[3] = select(neg_d, d, bool((byte >> (s + 3)) & 1));
|
||||
|
||||
reg = (type4) reg_f;
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q2_0(device const block_q2_0 * xb, short il, thread type4x4 & reg) {
|
||||
device const uint8_t * qs = xb->qs;
|
||||
const float d = xb->d;
|
||||
|
||||
const int byte_offset = il * 4; // il*16 elements = il*4 bytes (4 elements per byte)
|
||||
float4x4 reg_f;
|
||||
|
||||
for (int i = 0; i < 4; i++) {
|
||||
const uint8_t b = qs[byte_offset + i];
|
||||
reg_f[i][0] = ((float)((b >> 0) & 3) - 1.0f) * d;
|
||||
reg_f[i][1] = ((float)((b >> 2) & 3) - 1.0f) * d;
|
||||
reg_f[i][2] = ((float)((b >> 4) & 3) - 1.0f) * d;
|
||||
reg_f[i][3] = ((float)((b >> 6) & 3) - 1.0f) * d;
|
||||
}
|
||||
|
||||
reg = (type4x4) reg_f;
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_q2_0_t4(device const block_q2_0 * xb, short il, thread type4 & reg) {
|
||||
const float d = xb->d;
|
||||
const uint8_t b = xb->qs[il];
|
||||
|
||||
float4 reg_f;
|
||||
reg_f[0] = ((float)((b >> 0) & 3) - 1.0f) * d;
|
||||
reg_f[1] = ((float)((b >> 2) & 3) - 1.0f) * d;
|
||||
reg_f[2] = ((float)((b >> 4) & 3) - 1.0f) * d;
|
||||
reg_f[3] = ((float)((b >> 6) & 3) - 1.0f) * d;
|
||||
|
||||
reg = (type4) reg_f;
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q4_0(device const block_q4_0 * xb, short il, thread type4x4 & reg) {
|
||||
device const uint16_t * qs = ((device const uint16_t *)xb + 1);
|
||||
const float d1 = il ? (xb->d / 16.h) : xb->d;
|
||||
const float d2 = d1 / 256.f;
|
||||
const float md = -8.h * xb->d;
|
||||
const ushort mask0 = il ? 0x00F0 : 0x000F;
|
||||
const ushort mask1 = mask0 << 8;
|
||||
|
||||
float4x4 reg_f;
|
||||
|
||||
for (int i = 0; i < 8; i++) {
|
||||
reg_f[i/2][2*(i%2) + 0] = d1 * (qs[i] & mask0) + md;
|
||||
reg_f[i/2][2*(i%2) + 1] = d2 * (qs[i] & mask1) + md;
|
||||
}
|
||||
|
||||
reg = (type4x4) reg_f;
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_q4_0_t4(device const block_q4_0 * xb, short il, thread type4 & reg) {
|
||||
device const uint16_t * qs = ((device const uint16_t *)xb + 1);
|
||||
const float d1 = (il/4) ? (xb->d / 16.h) : xb->d;
|
||||
const float d2 = d1 / 256.f;
|
||||
const float md = -8.h * xb->d;
|
||||
const ushort mask0 = (il/4) ? 0x00F0 : 0x000F;
|
||||
const ushort mask1 = mask0 << 8;
|
||||
|
||||
for (int i = 0; i < 2; i++) {
|
||||
reg[2*i + 0] = d1 * (qs[2*(il%4) + i] & mask0) + md;
|
||||
reg[2*i + 1] = d2 * (qs[2*(il%4) + i] & mask1) + md;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q4_1(device const block_q4_1 * xb, short il, thread type4x4 & reg) {
|
||||
device const uint16_t * qs = ((device const uint16_t *)xb + 2);
|
||||
const float d1 = il ? (xb->d / 16.h) : xb->d;
|
||||
const float d2 = d1 / 256.f;
|
||||
const float m = xb->m;
|
||||
const ushort mask0 = il ? 0x00F0 : 0x000F;
|
||||
const ushort mask1 = mask0 << 8;
|
||||
|
||||
float4x4 reg_f;
|
||||
|
||||
for (int i = 0; i < 8; i++) {
|
||||
reg_f[i/2][2*(i%2) + 0] = ((qs[i] & mask0) * d1) + m;
|
||||
reg_f[i/2][2*(i%2) + 1] = ((qs[i] & mask1) * d2) + m;
|
||||
}
|
||||
|
||||
reg = (type4x4) reg_f;
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_q4_1_t4(device const block_q4_1 * xb, short il, thread type4 & reg) {
|
||||
device const uint16_t * qs = ((device const uint16_t *)xb + 2);
|
||||
const float d1 = (il/4) ? (xb->d / 16.h) : xb->d;
|
||||
const float d2 = d1 / 256.f;
|
||||
const float m = xb->m;
|
||||
const ushort mask0 = (il/4) ? 0x00F0 : 0x000F;
|
||||
const ushort mask1 = mask0 << 8;
|
||||
|
||||
for (int i = 0; i < 2; i++) {
|
||||
reg[2*i + 0] = d1 * (qs[2*(il%4) + i] & mask0) + m;
|
||||
reg[2*i + 1] = d2 * (qs[2*(il%4) + i] & mask1) + m;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q5_0(device const block_q5_0 * xb, short il, thread type4x4 & reg) {
|
||||
device const uint16_t * qs = ((device const uint16_t *)xb + 3);
|
||||
const float d = xb->d;
|
||||
const float md = -16.h * xb->d;
|
||||
const ushort mask = il ? 0x00F0 : 0x000F;
|
||||
|
||||
const uint32_t qh = *((device const uint32_t *)xb->qh);
|
||||
|
||||
const int x_mv = il ? 4 : 0;
|
||||
|
||||
const int gh_mv = il ? 12 : 0;
|
||||
const int gh_bk = il ? 0 : 4;
|
||||
|
||||
float4x4 reg_f;
|
||||
|
||||
for (int i = 0; i < 8; i++) {
|
||||
// extract the 5-th bits for x0 and x1
|
||||
const uint8_t xh_0 = ((qh >> (gh_mv + 2*i )) << gh_bk) & 0x10;
|
||||
const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
|
||||
|
||||
// combine the 4-bits from qs with the 5th bit
|
||||
const int32_t x0 = ((((qs[i] ) & mask) >> x_mv) | xh_0);
|
||||
const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
|
||||
|
||||
reg_f[i/2][2*(i%2) + 0] = d * x0 + md;
|
||||
reg_f[i/2][2*(i%2) + 1] = d * x1 + md;
|
||||
}
|
||||
|
||||
reg = (type4x4) reg_f;
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_q5_0_t4(device const block_q5_0 * xb, short il, thread type4 & reg) {
|
||||
device const uint16_t * qs = ((device const uint16_t *)xb + 3);
|
||||
const float d = xb->d;
|
||||
const float md = -16.h * xb->d;
|
||||
const ushort mask = (il/4) ? 0x00F0 : 0x000F;
|
||||
|
||||
const uint32_t qh = *((device const uint32_t *)xb->qh);
|
||||
|
||||
const int x_mv = (il/4) ? 4 : 0;
|
||||
|
||||
const int gh_mv = (il/4) ? 12 : 0;
|
||||
const int gh_bk = (il/4) ? 0 : 4;
|
||||
|
||||
for (int ii = 0; ii < 2; ii++) {
|
||||
int i = 2*(il%4) + ii;
|
||||
|
||||
// extract the 5-th bits for x0 and x1
|
||||
const uint8_t xh_0 = ((qh >> (gh_mv + 2*i )) << gh_bk) & 0x10;
|
||||
const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
|
||||
|
||||
// combine the 4-bits from qs with the 5th bit
|
||||
const int32_t x0 = ((((qs[i] ) & mask) >> x_mv) | xh_0);
|
||||
const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
|
||||
|
||||
reg[2*ii + 0] = d * x0 + md;
|
||||
reg[2*ii + 1] = d * x1 + md;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q5_1(device const block_q5_1 * xb, short il, thread type4x4 & reg) {
|
||||
device const uint16_t * qs = ((device const uint16_t *)xb + 4);
|
||||
const float d = xb->d;
|
||||
const float m = xb->m;
|
||||
const ushort mask = il ? 0x00F0 : 0x000F;
|
||||
|
||||
const uint32_t qh = *((device const uint32_t *)xb->qh);
|
||||
|
||||
const int x_mv = il ? 4 : 0;
|
||||
|
||||
const int gh_mv = il ? 12 : 0;
|
||||
const int gh_bk = il ? 0 : 4;
|
||||
|
||||
float4x4 reg_f;
|
||||
|
||||
for (int i = 0; i < 8; i++) {
|
||||
// extract the 5-th bits for x0 and x1
|
||||
const uint8_t xh_0 = ((qh >> (gh_mv + 2*i )) << gh_bk) & 0x10;
|
||||
const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
|
||||
|
||||
// combine the 4-bits from qs with the 5th bit
|
||||
const int32_t x0 = ((((qs[i] ) & mask) >> x_mv) | xh_0);
|
||||
const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
|
||||
|
||||
reg_f[i/2][2*(i%2) + 0] = d * x0 + m;
|
||||
reg_f[i/2][2*(i%2) + 1] = d * x1 + m;
|
||||
}
|
||||
|
||||
reg = (type4x4) reg_f;
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_q5_1_t4(device const block_q5_1 * xb, short il, thread type4 & reg) {
|
||||
device const uint16_t * qs = ((device const uint16_t *)xb + 4);
|
||||
const float d = xb->d;
|
||||
const float m = xb->m;
|
||||
const ushort mask = (il/4) ? 0x00F0 : 0x000F;
|
||||
|
||||
const uint32_t qh = *((device const uint32_t *)xb->qh);
|
||||
|
||||
const int x_mv = (il/4) ? 4 : 0;
|
||||
|
||||
const int gh_mv = (il/4) ? 12 : 0;
|
||||
const int gh_bk = (il/4) ? 0 : 4;
|
||||
|
||||
for (int ii = 0; ii < 2; ii++) {
|
||||
int i = 2*(il%4) + ii;
|
||||
|
||||
// extract the 5-th bits for x0 and x1
|
||||
const uint8_t xh_0 = ((qh >> (gh_mv + 2*i )) << gh_bk) & 0x10;
|
||||
const uint8_t xh_1 = ((qh >> (gh_mv + 2*i+1)) << gh_bk) & 0x10;
|
||||
|
||||
// combine the 4-bits from qs with the 5th bit
|
||||
const int32_t x0 = ((((qs[i] ) & mask) >> x_mv) | xh_0);
|
||||
const int32_t x1 = ((((qs[i] >> 8) & mask) >> x_mv) | xh_1);
|
||||
|
||||
reg[2*ii + 0] = d * x0 + m;
|
||||
reg[2*ii + 1] = d * x1 + m;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q8_0(device const block_q8_0 *xb, short il, thread type4x4 & reg) {
|
||||
device const int8_t * qs = ((device const int8_t *)xb->qs);
|
||||
const float d = xb->d;
|
||||
|
||||
float4x4 reg_f;
|
||||
|
||||
for (int i = 0; i < 16; i++) {
|
||||
reg_f[i/4][i%4] = (qs[i + 16*il] * d);
|
||||
}
|
||||
|
||||
reg = (type4x4) reg_f;
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_q8_0_t4(device const block_q8_0 *xb, short il, thread type4 & reg) {
|
||||
device const int8_t * qs = ((device const int8_t *)xb->qs);
|
||||
const float d = xb->d;
|
||||
|
||||
for (int i = 0; i < 4; i++) {
|
||||
reg[i] = (qs[4*(il%4) + i + 16*(il/4)] * d);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_mxfp4(device const block_mxfp4 * xb, short il, thread type4x4 & reg) {
|
||||
device const uint8_t * q2 = (device const uint8_t *)xb->qs;
|
||||
|
||||
const float d = e8m0_to_fp32(xb->e);
|
||||
const uint8_t shr = il >= 1 ? 4 : 0;
|
||||
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
reg[i][0] = d * kvalues_mxfp4_f[(q2[4*i + 0] >> shr) & 0x0F];
|
||||
reg[i][1] = d * kvalues_mxfp4_f[(q2[4*i + 1] >> shr) & 0x0F];
|
||||
reg[i][2] = d * kvalues_mxfp4_f[(q2[4*i + 2] >> shr) & 0x0F];
|
||||
reg[i][3] = d * kvalues_mxfp4_f[(q2[4*i + 3] >> shr) & 0x0F];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_mxfp4_t4(device const block_mxfp4 * xb, short il, thread type4 & reg) {
|
||||
device const uint8_t * q2 = (device const uint8_t *)xb->qs;
|
||||
|
||||
const float d = e8m0_to_fp32(xb->e);
|
||||
const short il4 = il%4;
|
||||
|
||||
const uint8_t shr = il >= 4 ? 4 : 0;
|
||||
|
||||
reg[0] = d * kvalues_mxfp4_f[(q2[4*il4 + 0] >> shr) & 0x0F];
|
||||
reg[1] = d * kvalues_mxfp4_f[(q2[4*il4 + 1] >> shr) & 0x0F];
|
||||
reg[2] = d * kvalues_mxfp4_f[(q2[4*il4 + 2] >> shr) & 0x0F];
|
||||
reg[3] = d * kvalues_mxfp4_f[(q2[4*il4 + 3] >> shr) & 0x0F];
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q2_K(device const block_q2_K *xb, short il, thread type4x4 & reg) {
|
||||
const float d = xb->d;
|
||||
const float min = xb->dmin;
|
||||
device const uint8_t * q = (device const uint8_t *)xb->qs;
|
||||
float dl, ml;
|
||||
uint8_t sc = xb->scales[il];
|
||||
|
||||
q = q + 32*(il/8) + 16*(il&1);
|
||||
il = (il/2)%4;
|
||||
|
||||
half coef = il>1 ? (il>2 ? 1/64.h : 1/16.h) : (il>0 ? 1/4.h : 1.h);
|
||||
uchar mask = il>1 ? (il>2 ? 192 : 48) : (il>0 ? 12 : 3);
|
||||
dl = d * (sc & 0xF) * coef, ml = min * (sc >> 4);
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
reg[i/4][i%4] = dl * (q[i] & mask) - ml;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q3_K(device const block_q3_K *xb, short il, thread type4x4 & reg) {
|
||||
const half d_all = xb->d;
|
||||
device const uint8_t * q = (device const uint8_t *)xb->qs;
|
||||
device const uint8_t * h = (device const uint8_t *)xb->hmask;
|
||||
device const int8_t * scales = (device const int8_t *)xb->scales;
|
||||
|
||||
q = q + 32 * (il/8) + 16 * (il&1);
|
||||
h = h + 16 * (il&1);
|
||||
uint8_t m = 1 << (il/2);
|
||||
uint16_t kmask1 = (il/4)>1 ? ((il/4)>2 ? 192 : 48) : \
|
||||
((il/4)>0 ? 12 : 3);
|
||||
uint16_t kmask2 = il/8 ? 0xF0 : 0x0F;
|
||||
uint16_t scale_2 = scales[il%8], scale_1 = scales[8 + il%4];
|
||||
int16_t dl_int = (il/4)&1 ? (scale_2&kmask2) | ((scale_1&kmask1) << 2)
|
||||
: (scale_2&kmask2) | ((scale_1&kmask1) << 4);
|
||||
float dl = il<8 ? d_all * (dl_int - 32.f) : d_all * (dl_int / 16.f - 32.f);
|
||||
const float ml = 4.f * dl;
|
||||
|
||||
il = (il/2) & 3;
|
||||
const half coef = il>1 ? (il>2 ? 1/64.h : 1/16.h) : (il>0 ? 1/4.h : 1.h);
|
||||
const uint8_t mask = il>1 ? (il>2 ? 192 : 48) : (il>0 ? 12 : 3);
|
||||
dl *= coef;
|
||||
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
reg[i/4][i%4] = dl * (q[i] & mask) - (h[i] & m ? 0 : ml);
|
||||
}
|
||||
}
|
||||
|
||||
static inline uchar2 get_scale_min_k4_just2(int j, int k, device const uchar * q) {
|
||||
return j < 4 ? uchar2{uchar(q[j+0+k] & 63), uchar(q[j+4+k] & 63)}
|
||||
: uchar2{uchar((q[j+4+k] & 0xF) | ((q[j-4+k] & 0xc0) >> 2)), uchar((q[j+4+k] >> 4) | ((q[j-0+k] & 0xc0) >> 2))};
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q4_K(device const block_q4_K * xb, short il, thread type4x4 & reg) {
|
||||
device const uchar * q = xb->qs;
|
||||
|
||||
short is = (il/4) * 2;
|
||||
q = q + (il/4) * 32 + 16 * (il&1);
|
||||
il = il & 3;
|
||||
const uchar2 sc = get_scale_min_k4_just2(is, il/2, xb->scales);
|
||||
const float d = il < 2 ? xb->d : xb->d / 16.h;
|
||||
const float min = xb->dmin;
|
||||
const float dl = d * sc[0];
|
||||
const float ml = min * sc[1];
|
||||
|
||||
const ushort mask = il < 2 ? 0x0F : 0xF0;
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
reg[i/4][i%4] = dl * (q[i] & mask) - ml;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q5_K(device const block_q5_K *xb, short il, thread type4x4 & reg) {
|
||||
device const uint8_t * q = xb->qs;
|
||||
device const uint8_t * qh = xb->qh;
|
||||
|
||||
short is = (il/4) * 2;
|
||||
q = q + 32 * (il/4) + 16 * (il&1);
|
||||
qh = qh + 16 * (il&1);
|
||||
uint8_t ul = 1 << (il/2);
|
||||
il = il & 3;
|
||||
const uchar2 sc = get_scale_min_k4_just2(is, il/2, xb->scales);
|
||||
const float d = il < 2 ? xb->d : xb->d / 16.f;
|
||||
const float min = xb->dmin;
|
||||
const float dl = d * sc[0];
|
||||
const float ml = min * sc[1];
|
||||
|
||||
const ushort mask = il<2 ? 0x0F : 0xF0;
|
||||
const float qh_val = il<2 ? 16.f : 256.f;
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
reg[i/4][i%4] = dl * ((q[i] & mask) + (qh[i] & ul ? qh_val : 0)) - ml;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_q6_K(device const block_q6_K *xb, short il, thread type4x4 & reg) {
|
||||
const half d_all = xb->d;
|
||||
device const uint16_t * ql = (device const uint16_t *)xb->ql;
|
||||
device const uint16_t * qh = (device const uint16_t *)xb->qh;
|
||||
device const int8_t * scales = (device const int8_t *)xb->scales;
|
||||
|
||||
ql = ql + 32*(il/8) + 16*((il/2)&1) + 8*(il&1);
|
||||
qh = qh + 16*(il/8) + 8*(il&1);
|
||||
float sc = scales[(il%2) + 2 * ((il/2))];
|
||||
il = (il/2) & 3;
|
||||
|
||||
const uint32_t kmask1 = il>1 ? (il>2 ? 0xC0C0C0C0 : 0x30303030) : (il>0 ? 0x0C0C0C0C : 0x03030303);
|
||||
const uint32_t kmask2 = il>1 ? 0xF0F0F0F0 : 0x0F0F0F0F;
|
||||
const float ml = d_all * sc * 32.f;
|
||||
const float dl0 = d_all * sc;
|
||||
const float dl1 = dl0 / 256.f;
|
||||
const float dl2 = dl0 / (256.f * 256.f);
|
||||
const float dl3 = dl0 / (256.f * 256.f * 256.f);
|
||||
const uint8_t shr_h = il>2 ? 2 : 0;
|
||||
const uint8_t shl_h = il>1 ? 0 : (il>0 ? 2 : 4);
|
||||
const uint8_t shr_l = il>1 ? 4 : 0;
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const uint32_t low = (ql[2*i] | (uint32_t)(ql[2*i+1] << 16)) & kmask2;
|
||||
const uint32_t high = (qh[2*i] | (uint32_t)(qh[2*i+1] << 16)) & kmask1;
|
||||
const uint32_t q = ((high << shl_h) >> shr_h) | (low >> shr_l);
|
||||
reg[i][0] = dl0 * ((half)(q & 0xFF)) - ml;
|
||||
reg[i][1] = dl1 * ((float)(q & 0xFF00)) - ml;
|
||||
reg[i][2] = dl2 * ((float)(q & 0xFF0000)) - ml;
|
||||
reg[i][3] = dl3 * ((float)(q & 0xFF000000)) - ml;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_iq2_xxs(device const block_iq2_xxs * xb, short il, thread type4x4 & reg) {
|
||||
// il is 0...15 for QK_K = 256 => index of block of 32 is il/2
|
||||
const float d = xb->d;
|
||||
const int ib32 = il/2;
|
||||
il = il%2;
|
||||
// il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
|
||||
// each block of 32 needs 2 uint32_t's for the quants & scale, so 4 uint16_t's.
|
||||
device const uint16_t * q2 = xb->qs + 4*ib32;
|
||||
const uint32_t aux32_g = q2[0] | (q2[1] << 16);
|
||||
const uint32_t aux32_s = q2[2] | (q2[3] << 16);
|
||||
thread const uint8_t * aux8 = (thread const uint8_t *)&aux32_g;
|
||||
const float dl = d * (0.5f + (aux32_s >> 28)) * 0.25f;
|
||||
constant uint8_t * grid = (constant uint8_t *)(iq2xxs_grid + aux8[2*il+0]);
|
||||
uint8_t signs = ksigns_iq2xs[(aux32_s >> 14*il) & 127];
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
reg[i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
|
||||
}
|
||||
grid = (constant uint8_t *)(iq2xxs_grid + aux8[2*il+1]);
|
||||
signs = ksigns_iq2xs[(aux32_s >> (14*il+7)) & 127];
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
reg[2+i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_iq2_xs(device const block_iq2_xs * xb, short il, thread type4x4 & reg) {
|
||||
// il is 0...15 for QK_K = 256 => index of block of 32 is il/2
|
||||
const float d = xb->d;
|
||||
const int ib32 = il/2;
|
||||
il = il%2;
|
||||
// il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
|
||||
device const uint16_t * q2 = xb->qs + 4*ib32;
|
||||
const float dl = d * (0.5f + ((xb->scales[ib32] >> 4*il) & 0xf)) * 0.25f;
|
||||
constant uint8_t * grid = (constant uint8_t *)(iq2xs_grid + (q2[2*il+0] & 511));
|
||||
uint8_t signs = ksigns_iq2xs[q2[2*il+0] >> 9];
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
reg[i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
|
||||
}
|
||||
grid = (constant uint8_t *)(iq2xs_grid + (q2[2*il+1] & 511));
|
||||
signs = ksigns_iq2xs[q2[2*il+1] >> 9];
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
reg[2+i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_iq3_xxs(device const block_iq3_xxs * xb, short il, thread type4x4 & reg) {
|
||||
// il is 0...15 for QK_K = 256 => index of block of 32 is il/2
|
||||
const float d = xb->d;
|
||||
const int ib32 = il/2;
|
||||
il = il%2;
|
||||
// il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
|
||||
device const uint8_t * q3 = xb->qs + 8*ib32;
|
||||
device const uint16_t * gas = (device const uint16_t *)(xb->qs + QK_K/4) + 2*ib32;
|
||||
const uint32_t aux32 = gas[0] | (gas[1] << 16);
|
||||
const float dl = d * (0.5f + (aux32 >> 28)) * 0.5f;
|
||||
constant uint8_t * grid1 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+0]);
|
||||
constant uint8_t * grid2 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+1]);
|
||||
uint8_t signs = ksigns_iq2xs[(aux32 >> 14*il) & 127];
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
reg[0][i] = dl * grid1[i] * (signs & kmask_iq2xs[i+0] ? -1.f : 1.f);
|
||||
reg[1][i] = dl * grid2[i] * (signs & kmask_iq2xs[i+4] ? -1.f : 1.f);
|
||||
}
|
||||
grid1 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+2]);
|
||||
grid2 = (constant uint8_t *)(iq3xxs_grid + q3[4*il+3]);
|
||||
signs = ksigns_iq2xs[(aux32 >> (14*il+7)) & 127];
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
reg[2][i] = dl * grid1[i] * (signs & kmask_iq2xs[i+0] ? -1.f : 1.f);
|
||||
reg[3][i] = dl * grid2[i] * (signs & kmask_iq2xs[i+4] ? -1.f : 1.f);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_iq3_s(device const block_iq3_s * xb, short il, thread type4x4 & reg) {
|
||||
// il is 0...15 for QK_K = 256 => index of block of 32 is il/2
|
||||
const float d = xb->d;
|
||||
const int ib32 = il/2;
|
||||
il = il%2;
|
||||
// il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
|
||||
device const uint8_t * qs = xb->qs + 8*ib32;
|
||||
device const uint8_t * signs = xb->signs + 4*ib32 + 2*il;
|
||||
const uint8_t qh = xb->qh[ib32] >> 4*il;
|
||||
const float dl = d * (1 + 2*((xb->scales[ib32/2] >> 4*(ib32%2)) & 0xf));
|
||||
constant uint8_t * grid1 = (constant uint8_t *)(iq3s_grid + (qs[4*il+0] | ((qh << 8) & 256)));
|
||||
constant uint8_t * grid2 = (constant uint8_t *)(iq3s_grid + (qs[4*il+1] | ((qh << 7) & 256)));
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
reg[0][i] = dl * grid1[i] * select(1, -1, signs[0] & kmask_iq2xs[i+0]);
|
||||
reg[1][i] = dl * grid2[i] * select(1, -1, signs[0] & kmask_iq2xs[i+4]);
|
||||
}
|
||||
grid1 = (constant uint8_t *)(iq3s_grid + (qs[4*il+2] | ((qh << 6) & 256)));
|
||||
grid2 = (constant uint8_t *)(iq3s_grid + (qs[4*il+3] | ((qh << 5) & 256)));
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
reg[2][i] = dl * grid1[i] * select(1, -1, signs[1] & kmask_iq2xs[i+0]);
|
||||
reg[3][i] = dl * grid2[i] * select(1, -1, signs[1] & kmask_iq2xs[i+4]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_iq2_s(device const block_iq2_s * xb, short il, thread type4x4 & reg) {
|
||||
// il is 0...15 for QK_K = 256 => index of block of 32 is il/2
|
||||
const float d = xb->d;
|
||||
const int ib32 = il/2;
|
||||
il = il%2;
|
||||
// il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
|
||||
device const uint8_t * qs = xb->qs + 4*ib32 + 2*il;
|
||||
device const uint8_t * signs = qs + QK_K/8;
|
||||
const uint8_t qh = xb->qh[ib32] >> 4*il;
|
||||
const float dl = d * (0.5f + ((xb->scales[ib32] >> 4*il) & 0xf)) * 0.25f;
|
||||
constant uint8_t * grid1 = (constant uint8_t *)(iq2s_grid + (qs[0] | ((qh << 8) & 0x300)));
|
||||
constant uint8_t * grid2 = (constant uint8_t *)(iq2s_grid + (qs[1] | ((qh << 6) & 0x300)));
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
reg[i/4+0][i%4] = dl * grid1[i] * select(1, -1, signs[0] & kmask_iq2xs[i]);
|
||||
reg[i/4+2][i%4] = dl * grid2[i] * select(1, -1, signs[1] & kmask_iq2xs[i]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_iq1_s(device const block_iq1_s * xb, short il, thread type4x4 & reg) {
|
||||
// il is 0...15 for QK_K = 256 => index of block of 32 is il/2
|
||||
const int ib32 = il/2;
|
||||
il = il%2;
|
||||
const float d = xb->d;
|
||||
device const uint8_t * qs = xb->qs + 4*ib32 + 2*il;
|
||||
device const uint16_t * qh = xb->qh;
|
||||
const float dl = d * (2*((qh[ib32] >> 12) & 7) + 1);
|
||||
const float ml = dl * (qh[ib32] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA);
|
||||
const uint16_t h = qh[ib32] >> 6*il;
|
||||
constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((h << 8) & 0x700)));
|
||||
constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((h << 5) & 0x700)));
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
reg[0][i] = dl * (grid1[i] & 0xf) + ml;
|
||||
reg[1][i] = dl * (grid1[i] >> 4) + ml;
|
||||
reg[2][i] = dl * (grid2[i] & 0xf) + ml;
|
||||
reg[3][i] = dl * (grid2[i] >> 4) + ml;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_iq1_m(device const block_iq1_m * xb, short il, thread type4x4 & reg) {
|
||||
// il is 0...15 for QK_K = 256 => index of block of 32 is il/2
|
||||
const int ib32 = il/2;
|
||||
il = il%2;
|
||||
device const uint16_t * sc = (device const uint16_t *)xb->scales;
|
||||
|
||||
iq1m_scale_t scale;
|
||||
scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
|
||||
const float d = scale.f16;
|
||||
|
||||
device const uint8_t * qs = xb->qs + 4*ib32 + 2*il;
|
||||
device const uint8_t * qh = xb->qh + 2*ib32 + il;
|
||||
|
||||
const float dl = d * (2*((sc[ib32/2] >> (6*(ib32%2)+3*il)) & 7) + 1);
|
||||
const float ml1 = dl * (qh[0] & 0x08 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA);
|
||||
const float ml2 = dl * (qh[0] & 0x80 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA);
|
||||
constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
|
||||
constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 4) & 0x700)));
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
reg[0][i] = dl * (grid1[i] & 0xf) + ml1;
|
||||
reg[1][i] = dl * (grid1[i] >> 4) + ml1;
|
||||
reg[2][i] = dl * (grid2[i] & 0xf) + ml2;
|
||||
reg[3][i] = dl * (grid2[i] >> 4) + ml2;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_iq4_nl(device const block_iq4_nl * xb, short il, thread type4x4 & reg) {
|
||||
device const uint16_t * q4 = (device const uint16_t *)xb->qs;
|
||||
const float d = xb->d;
|
||||
uint32_t aux32;
|
||||
thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
aux32 = ((q4[2*i] | (q4[2*i+1] << 16)) >> 4*il) & 0x0f0f0f0f;
|
||||
reg[i][0] = d * kvalues_iq4nl_f[q8[0]];
|
||||
reg[i][1] = d * kvalues_iq4nl_f[q8[1]];
|
||||
reg[i][2] = d * kvalues_iq4nl_f[q8[2]];
|
||||
reg[i][3] = d * kvalues_iq4nl_f[q8[3]];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename type4>
|
||||
void dequantize_iq4_nl_t4(device const block_iq4_nl * xb, short il, thread type4 & reg) {
|
||||
device const uint16_t * q4 = (device const uint16_t *)xb->qs;
|
||||
const float d = xb->d;
|
||||
uint32_t aux32;
|
||||
thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
|
||||
aux32 = ((q4[2*(il%4)] | (q4[2*(il%4)+1] << 16)) >> 4*(il/4)) & 0x0f0f0f0f;
|
||||
reg[0] = d * kvalues_iq4nl_f[q8[0]];
|
||||
reg[1] = d * kvalues_iq4nl_f[q8[1]];
|
||||
reg[2] = d * kvalues_iq4nl_f[q8[2]];
|
||||
reg[3] = d * kvalues_iq4nl_f[q8[3]];
|
||||
}
|
||||
|
||||
template <typename type4x4>
|
||||
void dequantize_iq4_xs(device const block_iq4_xs * xb, short il, thread type4x4 & reg) {
|
||||
// il is 0...15 for QK_K = 256 => index of block of 32 is il/2
|
||||
const int ib32 = il/2;
|
||||
il = il%2;
|
||||
// il = 0 or 1. il = 0 processes the first 16 quants in a block of 32, il = 1 the second 16
|
||||
device const uint32_t * q4 = (device const uint32_t *)xb->qs + 4*ib32;
|
||||
const int ls = ((xb->scales_l[ib32/2] >> 4*(ib32%2)) & 0xf) | (((xb->scales_h >> 2*ib32) & 3) << 4);
|
||||
const float d = (float)xb->d * (ls - 32);
|
||||
uint32_t aux32;
|
||||
thread const uint8_t * q8 = (thread const uint8_t *)&aux32;
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
aux32 = (q4[i] >> 4*il) & 0x0f0f0f0f;
|
||||
reg[i][0] = d * kvalues_iq4nl_f[q8[0]];
|
||||
reg[i][1] = d * kvalues_iq4nl_f[q8[1]];
|
||||
reg[i][2] = d * kvalues_iq4nl_f[q8[2]];
|
||||
reg[i][3] = d * kvalues_iq4nl_f[q8[3]];
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,250 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
constant short FC_gated_delta_net_ne20 [[function_constant(FC_GATED_DELTA_NET + 0)]];
|
||||
constant short FC_gated_delta_net_ne30 [[function_constant(FC_GATED_DELTA_NET + 1)]];
|
||||
constant short FC_gated_delta_net_K [[function_constant(FC_GATED_DELTA_NET + 2)]];
|
||||
|
||||
#if 1
|
||||
template<short NSG>
|
||||
kernel void kernel_gated_delta_net_impl(
|
||||
constant ggml_metal_kargs_gated_delta_net & args,
|
||||
device const char * q,
|
||||
device const char * k,
|
||||
device const char * v,
|
||||
device const char * g,
|
||||
device const char * b,
|
||||
device const char * s,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
#define S_v FC_gated_delta_net_ne20
|
||||
#define G FC_gated_delta_net_ne30
|
||||
#define K FC_gated_delta_net_K
|
||||
|
||||
const uint tx = tpitg.x;
|
||||
const uint ty = tpitg.y;
|
||||
|
||||
const uint i23 = tgpig.z; // B (n_seqs)
|
||||
const uint i21 = tgpig.y; // H (head)
|
||||
const uint i20 = tgpig.x*NSG + ty; // row within S_v
|
||||
|
||||
const uint i01 = i21 % args.ne01;
|
||||
const uint i11 = i21 % args.ne11;
|
||||
|
||||
const float scale = 1.0f / sqrt((float)S_v);
|
||||
|
||||
// input state layout [S_v, S_v, H, n_seqs] (s0 only): per-seq stride is H*D.
|
||||
// state is stored transposed: M[i20][is] = S[is][i20], so row i20 is contiguous
|
||||
const uint state_in_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
|
||||
device const float * s_ptr = (device const float *) (s) + state_in_base;
|
||||
|
||||
float ls[NSG];
|
||||
|
||||
FOR_UNROLL (short j = 0; j < NSG; j++) {
|
||||
const short is = tx*NSG + j;
|
||||
ls[j] = s_ptr[is];
|
||||
}
|
||||
|
||||
device float * dst_attn = (device float *) (dst) + (i23*args.ne22*args.ne21 + i21)*S_v + i20;
|
||||
|
||||
device const float * q_ptr = (device const float *) (q + i23*args.nb03 + i01*args.nb01);
|
||||
device const float * k_ptr = (device const float *) (k + i23*args.nb13 + i11*args.nb11);
|
||||
device const float * v_ptr = (device const float *) (v + i23*args.nb23 + i21*args.nb21);
|
||||
|
||||
device const float * b_ptr = (device const float *) (b) + (i23*args.ne22*args.ne21 + i21);
|
||||
device const float * g_ptr = (device const float *) (g) + (i23*args.ne22*args.ne21 + i21)*G;
|
||||
|
||||
// snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
|
||||
// When n_tokens < K, only slots 0..n_tokens-1 are written; older slots are caller-owned.
|
||||
|
||||
// output state base offset: after attention scores
|
||||
const uint attn_size = args.ne22 * args.ne21 * S_v * args.ne23;
|
||||
// output state per-slot size: S_v * S_v * H * n_seqs
|
||||
const uint state_size_per_snap = S_v * S_v * args.ne21 * args.ne23;
|
||||
// per-(seq,head) offset within a slot
|
||||
const uint state_out_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
|
||||
|
||||
for (short t = 0; t < args.ne22; t++) {
|
||||
float s_k = 0.0f;
|
||||
|
||||
if (G == 1) {
|
||||
const float g_exp = exp(g_ptr[0]);
|
||||
|
||||
FOR_UNROLL (short j = 0; j < NSG; j++) {
|
||||
const short is = tx*NSG + j;
|
||||
ls[j] *= g_exp;
|
||||
|
||||
s_k += ls[j]*k_ptr[is];
|
||||
}
|
||||
} else {
|
||||
// KDA
|
||||
FOR_UNROLL (short j = 0; j < NSG; j++) {
|
||||
const short is = tx*NSG + j;
|
||||
ls[j] *= exp(g_ptr[is]);
|
||||
|
||||
s_k += ls[j]*k_ptr[is];
|
||||
}
|
||||
}
|
||||
|
||||
s_k = simd_sum(s_k);
|
||||
|
||||
const float d = (v_ptr[i20] - s_k)*b_ptr[0];
|
||||
|
||||
float y = 0.0f;
|
||||
|
||||
FOR_UNROLL (short j = 0; j < NSG; j++) {
|
||||
const short is = tx*NSG + j;
|
||||
ls[j] += k_ptr[is]*d;
|
||||
|
||||
y += ls[j]*q_ptr[is];
|
||||
}
|
||||
|
||||
y = simd_sum(y);
|
||||
|
||||
if (tx == 0) {
|
||||
dst_attn[t*args.ne21*S_v] = y*scale;
|
||||
}
|
||||
|
||||
q_ptr += args.ns02;
|
||||
k_ptr += args.ns12;
|
||||
v_ptr += args.ns22;
|
||||
|
||||
b_ptr += args.ne21;
|
||||
g_ptr += args.ne21*G;
|
||||
|
||||
if (K > 1) {
|
||||
const int target_slot = (int)args.ne22 - 1 - (int)t;
|
||||
if (target_slot >= 0 && target_slot < (int)K) {
|
||||
device float * dst_state = (device float *) (dst) + attn_size + (uint)target_slot * state_size_per_snap + state_out_base;
|
||||
FOR_UNROLL (short j = 0; j < NSG; j++) {
|
||||
const short is = tx*NSG + j;
|
||||
dst_state[is] = ls[j];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (K == 1) {
|
||||
device float * dst_state = (device float *) (dst) + attn_size + state_out_base;
|
||||
FOR_UNROLL (short j = 0; j < NSG; j++) {
|
||||
const short is = tx*NSG + j;
|
||||
dst_state[is] = ls[j];
|
||||
}
|
||||
}
|
||||
|
||||
#undef S_v
|
||||
#undef G
|
||||
#undef K
|
||||
}
|
||||
|
||||
typedef decltype(kernel_gated_delta_net_impl<4>) kernel_gated_delta_net_t;
|
||||
|
||||
template [[host_name("kernel_gated_delta_net_f32_1")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<1>;
|
||||
template [[host_name("kernel_gated_delta_net_f32_2")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<2>;
|
||||
template [[host_name("kernel_gated_delta_net_f32_4")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<4>;
|
||||
|
||||
#else
|
||||
// a simplified version of the above
|
||||
// no performance improvement, so keep the above version for now
|
||||
|
||||
template<typename T, short NSG>
|
||||
kernel void kernel_gated_delta_net_impl(
|
||||
constant ggml_metal_kargs_gated_delta_net & args,
|
||||
device const char * q,
|
||||
device const char * k,
|
||||
device const char * v,
|
||||
device const char * g,
|
||||
device const char * b,
|
||||
device const char * s,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
#define S_v FC_gated_delta_net_ne20
|
||||
#define G FC_gated_delta_net_ne30
|
||||
|
||||
const uint tx = tpitg.x;
|
||||
const uint ty = tpitg.y;
|
||||
|
||||
const uint i23 = tgpig.z; // B
|
||||
const uint i21 = tgpig.y; // H
|
||||
const uint i20 = tgpig.x*NSG + ty;
|
||||
|
||||
const uint i01 = i21 % args.ne01;
|
||||
const uint i11 = i21 % args.ne11;
|
||||
|
||||
const float scale = 1.0f / sqrt((float)S_v);
|
||||
|
||||
device const float * s_ptr = (device const float *) (s) + (i23*args.ne21 + i21)*S_v*S_v + i20;
|
||||
|
||||
float lsf[NSG];
|
||||
|
||||
FOR_UNROLL (short j = 0; j < NSG; j++) {
|
||||
const short is = tx*NSG + j;
|
||||
lsf[j] = s_ptr[is*S_v];
|
||||
}
|
||||
|
||||
thread T * ls = (thread T *) (lsf);
|
||||
|
||||
device float * dst_attn = (device float *) (dst) + (i23*args.ne22*args.ne21 + i21)*S_v + i20;
|
||||
|
||||
device const float * q_ptr = (device const float *) (q + i23*args.nb03 + i01*args.nb01);
|
||||
device const float * k_ptr = (device const float *) (k + i23*args.nb13 + i11*args.nb11);
|
||||
device const float * v_ptr = (device const float *) (v + i23*args.nb23 + i21*args.nb21);
|
||||
|
||||
device const float * b_ptr = (device const float *) (b) + (i23*args.ne22*args.ne21 + i21);
|
||||
device const float * g_ptr = (device const float *) (g) + (i23*args.ne22*args.ne21 + i21)*G;
|
||||
|
||||
for (short t = 0; t < args.ne22; t++) {
|
||||
device const T * qt_ptr = (device const T *) (q_ptr);
|
||||
device const T * kt_ptr = (device const T *) (k_ptr);
|
||||
device const T * gt_ptr = (device const T *) (g_ptr);
|
||||
|
||||
if (G == 1) {
|
||||
*ls *= exp(g_ptr[0]);
|
||||
} else {
|
||||
// KDA
|
||||
*ls *= exp(gt_ptr[tx]);
|
||||
}
|
||||
|
||||
const float s_k = simd_sum(dot(*ls, kt_ptr[tx]));
|
||||
|
||||
const float d = (v_ptr[i20] - s_k)*b_ptr[0];
|
||||
|
||||
*ls += kt_ptr[tx]*d;
|
||||
|
||||
const float y = simd_sum(dot(*ls, qt_ptr[tx]));
|
||||
|
||||
if (tx == 0) {
|
||||
*dst_attn = y*scale;
|
||||
}
|
||||
|
||||
q_ptr += args.ns02;
|
||||
k_ptr += args.ns12;
|
||||
v_ptr += args.ns22;
|
||||
|
||||
b_ptr += args.ne21;
|
||||
g_ptr += args.ne21*G;
|
||||
|
||||
dst_attn += args.ne21*S_v;
|
||||
}
|
||||
|
||||
device float * dst_state = (device float *) (dst) + args.ne23*args.ne22*args.ne21*S_v + (i23*args.ne21 + i21)*S_v*S_v + i20;
|
||||
device T * dstt_state = (device T *) (dst_state);
|
||||
|
||||
FOR_UNROLL (short j = 0; j < NSG; j++) {
|
||||
const short is = tx*NSG + j;
|
||||
dst_state[is*S_v] = lsf[j];
|
||||
}
|
||||
|
||||
#undef S_v
|
||||
#undef G
|
||||
}
|
||||
|
||||
typedef decltype(kernel_gated_delta_net_impl<float4, 4>) kernel_gated_delta_net_t;
|
||||
|
||||
template [[host_name("kernel_gated_delta_net_f32_1")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float, 1>;
|
||||
template [[host_name("kernel_gated_delta_net_f32_2")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float2, 2>;
|
||||
template [[host_name("kernel_gated_delta_net_f32_4")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float4, 4>;
|
||||
#endif
|
||||
@@ -1,595 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
kernel void kernel_argmax_f32(
|
||||
constant ggml_metal_kargs_argmax & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
uint tpitg[[thread_position_in_threadgroup]],
|
||||
uint sgitg[[simdgroup_index_in_threadgroup]],
|
||||
uint tiisg[[thread_index_in_simdgroup]],
|
||||
uint ntg[[threads_per_threadgroup]]) {
|
||||
device const float * x_row = (device const float *) ((device const char *) src0 + tgpig * args.nb01);
|
||||
|
||||
float lmax = -INFINITY;
|
||||
int32_t larg = -1;
|
||||
|
||||
for (int i00 = tpitg; i00 < args.ne00; i00 += ntg) {
|
||||
if (x_row[i00] > lmax) {
|
||||
lmax = x_row[i00];
|
||||
larg = i00;
|
||||
}
|
||||
}
|
||||
|
||||
// find the argmax value in the block
|
||||
float max_val = simd_max(lmax);
|
||||
int32_t arg_val = simd_max(select(-1, larg, lmax == max_val));
|
||||
|
||||
device int32_t * dst_i32 = (device int32_t *) dst;
|
||||
|
||||
threadgroup float * shared_maxval = (threadgroup float *) shmem;
|
||||
threadgroup int32_t * shared_argmax = (threadgroup int32_t *) shmem + N_SIMDWIDTH;
|
||||
|
||||
if (ntg > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
shared_maxval[tiisg] = -INFINITY;
|
||||
shared_argmax[tiisg] = -1;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shared_maxval[sgitg] = max_val;
|
||||
shared_argmax[sgitg] = arg_val;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
max_val = shared_maxval[tiisg];
|
||||
arg_val = shared_argmax[tiisg];
|
||||
|
||||
float max_val_reduced = simd_max(max_val);
|
||||
int32_t arg_val_reduced = simd_max(select(-1, arg_val, max_val == max_val_reduced));
|
||||
|
||||
dst_i32[tgpig] = arg_val_reduced;
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
dst_i32[tgpig] = arg_val;
|
||||
}
|
||||
|
||||
kernel void kernel_diag_f32(
|
||||
constant ggml_metal_kargs_diag & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]]) {
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
|
||||
const int32_t i3 = tgpig.z;
|
||||
const int32_t i2 = tgpig.y;
|
||||
const int32_t i1 = tgpig.x;
|
||||
|
||||
device const float * src0_ptr = (device const float *)(src0 + i2*args.nb02 + i3*args.nb03);
|
||||
device float * dst_ptr = (device float *)(dst + i1*args.nb01 + i2*args.nb2 + i3*args.nb3);
|
||||
|
||||
for (int i0 = tiitg; i0 < args.ne0; i0 += NW) {
|
||||
dst_ptr[i0] = i0 == i1 ? src0_ptr[i0] : 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_roll_f32(
|
||||
constant ggml_metal_kargs_roll & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int64_t i3 = tgpig.z;
|
||||
const int64_t i2 = tgpig.y;
|
||||
const int64_t i1 = tgpig.x;
|
||||
|
||||
device const float * src0_ptr = (device const float *) src0;
|
||||
device float * dst_ptr = (device float *) dst;
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
// apply shifts and wrap around
|
||||
int64_t i00 = i0 - args.s0;
|
||||
int64_t i01 = i1 - args.s1;
|
||||
int64_t i02 = i2 - args.s2;
|
||||
int64_t i03 = i3 - args.s3;
|
||||
|
||||
if (i00 < 0) { i00 += args.ne00; } else if (i00 >= args.ne00) { i00 -= args.ne00; }
|
||||
if (i01 < 0) { i01 += args.ne01; } else if (i01 >= args.ne01) { i01 -= args.ne01; }
|
||||
if (i02 < 0) { i02 += args.ne02; } else if (i02 >= args.ne02) { i02 -= args.ne02; }
|
||||
if (i03 < 0) { i03 += args.ne03; } else if (i03 >= args.ne03) { i03 -= args.ne03; }
|
||||
|
||||
int64_t src_idx = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00 + i00;
|
||||
int64_t dst_idx = i3 *args.ne2 *args.ne1 *args.ne0 + i2 *args.ne1 *args.ne0 + i1 *args.ne0 + i0;
|
||||
|
||||
dst_ptr[dst_idx] = src0_ptr[src_idx];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
kernel void kernel_pad_impl(
|
||||
constant ggml_metal_kargs_pad & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
const int32_t i3 = tgpig.z;
|
||||
const int32_t i2 = tgpig.y;
|
||||
const int32_t k0 = tgpig.x/args.ne1;
|
||||
const int32_t i1 = tgpig.x - k0*args.ne1;
|
||||
|
||||
const int32_t i03 = i3;
|
||||
const int32_t i02 = i2;
|
||||
const int32_t i01 = i1;
|
||||
|
||||
device const T * src0_ptr = (device const T *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
|
||||
device T * dst_ptr = (device T *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1);
|
||||
|
||||
for (int32_t l0 = 0; l0 < 1024; l0 += ntg.x) {
|
||||
const int32_t i0 = k0*1024 + tpitg.x + l0;
|
||||
if (i0 >= args.ne0) {
|
||||
break;
|
||||
}
|
||||
|
||||
if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
|
||||
dst_ptr[i0] = src0_ptr[i0];
|
||||
} else {
|
||||
dst_ptr[i0] = 0.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_pad_impl<float>) kernel_pad_t;
|
||||
|
||||
template [[host_name("kernel_pad_f32")]] kernel kernel_pad_t kernel_pad_impl<float>;
|
||||
template [[host_name("kernel_pad_f32_4")]] kernel kernel_pad_t kernel_pad_impl<float4>;
|
||||
|
||||
// TODO: this is slow - optimize
|
||||
kernel void kernel_pad_reflect_1d_f32(
|
||||
constant ggml_metal_kargs_pad_reflect_1d & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tgpg[[threadgroups_per_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int64_t i3 = tgpig.z;
|
||||
const int64_t i2 = tgpig.y;
|
||||
const int64_t i1 = tgpig.x;
|
||||
|
||||
const int64_t i03 = i3;
|
||||
const int64_t i02 = i2;
|
||||
const int64_t i01 = i1;
|
||||
|
||||
device const float * src0_ptr = (device const float *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
|
||||
device float * dst_ptr = (device float *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1);
|
||||
|
||||
if (i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
if (i0 < args.p0) {
|
||||
dst_ptr[i0] = src0_ptr[args.p0 - i0];
|
||||
} else if (i0 < args.ne0 - args.p1) {
|
||||
dst_ptr[i0] = src0_ptr[i0 - args.p0];
|
||||
} else {
|
||||
dst_ptr[i0] = src0_ptr[(args.ne0 - args.p1 - args.p0) - (args.p1 + 1 - (args.ne0 - i0)) - 1];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_arange_f32(
|
||||
constant ggml_metal_kargs_arange & args,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
device float * dst_ptr = (device float *) dst;
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
dst_ptr[i0] = args.start + args.step * i0;
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_timestep_embedding_f32(
|
||||
constant ggml_metal_kargs_timestep_embedding & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
int i = tgpig.x;
|
||||
device float * embed_data = (device float *)(dst + i*args.nb1);
|
||||
|
||||
int half_ = args.dim / 2;
|
||||
for (int j = tpitg.x; j < half_; j += ntg.x) {
|
||||
float timestep = ((device float *)src0)[i];
|
||||
float freq = (float)exp(-log((float)args.max_period) * j / half_);
|
||||
float arg = timestep * freq;
|
||||
embed_data[j ] = cos(arg);
|
||||
embed_data[j + half_] = sin(arg);
|
||||
}
|
||||
|
||||
if (args.dim % 2 != 0 && tpitg.x == 0) {
|
||||
embed_data[2 * half_] = 0.f;
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_opt_step_adamw_f32(
|
||||
constant ggml_metal_kargs_opt_step_adamw & args,
|
||||
device float * x,
|
||||
device const float * g,
|
||||
device float * g_m,
|
||||
device float * g_v,
|
||||
device const float * pars,
|
||||
uint gid[[thread_position_in_grid]]) {
|
||||
|
||||
if (gid >= args.np) {
|
||||
return;
|
||||
}
|
||||
|
||||
const float alpha = pars[0];
|
||||
const float beta1 = pars[1];
|
||||
const float beta2 = pars[2];
|
||||
const float eps = pars[3];
|
||||
const float wd = pars[4];
|
||||
const float beta1h = pars[5];
|
||||
const float beta2h = pars[6];
|
||||
|
||||
const float gi = g[gid];
|
||||
const float gmi = g_m[gid] * beta1 + gi * (1.0f - beta1);
|
||||
const float gvi = g_v[gid] * beta2 + gi * gi * (1.0f - beta2);
|
||||
|
||||
g_m[gid] = gmi;
|
||||
g_v[gid] = gvi;
|
||||
|
||||
const float mh = gmi * beta1h;
|
||||
const float vh = sqrt(gvi * beta2h) + eps;
|
||||
|
||||
x[gid] = x[gid] * (1.0f - alpha * wd) - alpha * mh / vh;
|
||||
}
|
||||
|
||||
kernel void kernel_opt_step_sgd_f32(
|
||||
constant ggml_metal_kargs_opt_step_sgd & args,
|
||||
device float * x,
|
||||
device const float * g,
|
||||
device const float * pars,
|
||||
uint gid[[thread_position_in_grid]]) {
|
||||
|
||||
if (gid >= args.np) {
|
||||
return;
|
||||
}
|
||||
|
||||
x[gid] = x[gid] * (1.0f - pars[0] * pars[1]) - pars[0] * g[gid];
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_memset(
|
||||
constant ggml_metal_kargs_memset & args,
|
||||
device T * dst,
|
||||
uint tpig[[thread_position_in_grid]]) {
|
||||
dst[tpig] = args.val;
|
||||
}
|
||||
|
||||
typedef decltype(kernel_memset<int64_t>) kernel_memset_t;
|
||||
|
||||
template [[host_name("kernel_memset_i64")]] kernel kernel_memset_t kernel_memset<int64_t>;
|
||||
|
||||
constant short FC_count_equal_nsg [[function_constant(FC_COUNT_EQUAL + 0)]];
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_count_equal(
|
||||
constant ggml_metal_kargs_count_equal & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device atomic_int * dst,
|
||||
threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const short NSG = FC_count_equal_nsg;
|
||||
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = tgpig.x;
|
||||
|
||||
if (i3 >= args.ne03 || i2 >= args.ne02 || i1 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
int sum = 0;
|
||||
|
||||
device const char * base0 = src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03;
|
||||
device const char * base1 = src1 + i1*args.nb11 + i2*args.nb12 + i3*args.nb13;
|
||||
|
||||
for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
|
||||
const T v0 = *(device const T *)(base0 + i0*args.nb00);
|
||||
const T v1 = *(device const T *)(base1 + i0*args.nb10);
|
||||
sum += (v0 == v1);
|
||||
}
|
||||
|
||||
sum = simd_sum(sum);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_i32[sgitg] = sum;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (sgitg == 0) {
|
||||
float v = 0.0f;
|
||||
if (tpitg.x < NSG) {
|
||||
v = shmem_i32[tpitg.x];
|
||||
}
|
||||
|
||||
float total = simd_sum(v);
|
||||
if (tpitg.x == 0) {
|
||||
atomic_fetch_add_explicit(dst, (int32_t) total, memory_order_relaxed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_count_equal<int32_t>) kernel_count_equal_t;
|
||||
|
||||
template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal<int32_t>;
|
||||
|
||||
template <typename T>
|
||||
kernel void kernel_snake(
|
||||
constant ggml_metal_kargs_snake & args,
|
||||
device const T * x,
|
||||
device const float * a,
|
||||
device const float * inv_b,
|
||||
device T * dst,
|
||||
uint tgpig [[threadgroup_position_in_grid]],
|
||||
uint tpitg [[thread_position_in_threadgroup]],
|
||||
uint ntg [[threads_per_threadgroup]]) {
|
||||
|
||||
const int idx = tgpig * ntg + tpitg;
|
||||
if (idx >= args.T * args.C) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int c = idx / args.T; // x is [T, C], a / inv_b collapse to [1, C]
|
||||
const float xi = float(x[idx]);
|
||||
const float si = sin(a[c] * xi);
|
||||
dst[idx] = T(xi + si * si * inv_b[c]);
|
||||
}
|
||||
|
||||
template [[host_name("kernel_snake_f32")]] kernel void kernel_snake<float>(constant ggml_metal_kargs_snake &, device const float *, device const float *, device const float *, device float *, uint, uint, uint);
|
||||
template [[host_name("kernel_snake_f16")]] kernel void kernel_snake<half>(constant ggml_metal_kargs_snake &, device const half *, device const float *, device const float *, device half *, uint, uint, uint);
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_snake_bf16")]] kernel void kernel_snake<bfloat>(constant ggml_metal_kargs_snake &, device const bfloat *, device const float *, device const float *, device bfloat *, uint, uint, uint);
|
||||
#endif
|
||||
|
||||
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>;
|
||||
|
||||
kernel void kernel_dsv4_hc_comb_f32(
|
||||
constant ggml_metal_kargs_dsv4_hc_comb & args,
|
||||
device const char * mixes,
|
||||
device const char * scale,
|
||||
device const char * base,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr ushort hc = 4;
|
||||
constexpr ushort comb_offset = 2*hc;
|
||||
|
||||
const int it = tgpig.x*ntg.y + sgitg;
|
||||
if (it >= args.n_tokens) {
|
||||
return;
|
||||
}
|
||||
|
||||
float scale_lane = 0.0f;
|
||||
if (tiisg == 0) {
|
||||
scale_lane = *(device const float *) (scale + 2*args.nb_s0);
|
||||
}
|
||||
const float scale_comb = simd_shuffle(scale_lane, 0);
|
||||
|
||||
float v = 0.0f;
|
||||
if (tiisg < hc*hc) {
|
||||
v = *(device const float *) (mixes + (comb_offset + tiisg)*args.nb_m0 + it*args.nb_m1)*scale_comb
|
||||
+ *(device const float *) (base + (comb_offset + tiisg)*args.nb_b0);
|
||||
}
|
||||
|
||||
// Softmax across destinations (the four contiguous lanes for each source).
|
||||
float vmax = max(v, simd_shuffle_xor(v, 1));
|
||||
vmax = max(vmax, simd_shuffle_xor(vmax, 2));
|
||||
v = exp(v - vmax);
|
||||
|
||||
float sum = v + simd_shuffle_xor(v, 1);
|
||||
sum += simd_shuffle_xor(sum, 2);
|
||||
v = v/sum + args.eps;
|
||||
|
||||
// Normalize columns: equal destination indices are four lanes apart.
|
||||
sum = v + simd_shuffle_xor(v, 4);
|
||||
sum += simd_shuffle_xor(sum, 8);
|
||||
v /= sum + args.eps;
|
||||
|
||||
for (int i = 1; i < args.n_iter; ++i) {
|
||||
sum = v + simd_shuffle_xor(v, 1);
|
||||
sum += simd_shuffle_xor(sum, 2);
|
||||
v /= sum + args.eps;
|
||||
|
||||
sum = v + simd_shuffle_xor(v, 4);
|
||||
sum += simd_shuffle_xor(sum, 8);
|
||||
v /= sum + args.eps;
|
||||
}
|
||||
|
||||
if (tiisg < hc*hc) {
|
||||
const ushort idst = tiisg & 3;
|
||||
const ushort isrc = tiisg >> 2;
|
||||
*(device float *) (dst + idst*args.nb_d0 + isrc*args.nb_d1 + it*args.nb_d2) = v;
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_hc_pre_f32(
|
||||
constant ggml_metal_kargs_dsv4_hc_pre & args,
|
||||
device const char * x,
|
||||
device const char * weights,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr ushort hc = 4;
|
||||
|
||||
const int it = tgpig.y;
|
||||
const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg;
|
||||
|
||||
float weight_lane = 0.0f;
|
||||
if (tiisg < hc) {
|
||||
weight_lane = *(device const float *) (weights + tiisg*args.nb_w0 + it*args.nb_w1);
|
||||
}
|
||||
|
||||
float w[hc];
|
||||
FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) {
|
||||
w[ih] = simd_shuffle(weight_lane, ih);
|
||||
}
|
||||
|
||||
if (i0 >= args.n_embd) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const char * xb = x + i0*args.nb_x0 + it*args.nb_x2;
|
||||
float result = 0.0f;
|
||||
FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) {
|
||||
result = fma(*(device const float *) (xb + ih*args.nb_x1), w[ih], result);
|
||||
}
|
||||
|
||||
*(device float *) (dst + i0*args.nb_d0 + it*args.nb_d1) = result;
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_hc_post_f32(
|
||||
constant ggml_metal_kargs_dsv4_hc_post & args,
|
||||
device const char * x,
|
||||
device const char * residual,
|
||||
device const char * post,
|
||||
device const char * comb,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr ushort hc = 4;
|
||||
|
||||
const int it = tgpig.y;
|
||||
const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg;
|
||||
|
||||
float coeff_lane = 0.0f;
|
||||
if (tiisg < hc) {
|
||||
coeff_lane = *(device const float *) (post + tiisg*args.nb_p0 + it*args.nb_p1);
|
||||
} else if (tiisg < hc + hc*hc) {
|
||||
const ushort idx = tiisg - hc;
|
||||
const ushort idst = idx & 3;
|
||||
const ushort isrc = idx >> 2;
|
||||
coeff_lane = *(device const float *) (comb + idst*args.nb_c0 + isrc*args.nb_c1 + it*args.nb_c2);
|
||||
}
|
||||
|
||||
float post_reg[hc];
|
||||
float comb_reg[hc][hc];
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
post_reg[idst] = simd_shuffle(coeff_lane, idst);
|
||||
}
|
||||
FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) {
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
comb_reg[isrc][idst] = simd_shuffle(coeff_lane, hc + idst + hc*isrc);
|
||||
}
|
||||
}
|
||||
|
||||
if (i0 >= args.n_embd) {
|
||||
return;
|
||||
}
|
||||
|
||||
const float xv = *(device const float *) (x + i0*args.nb_x0 + it*args.nb_x1);
|
||||
float result[hc];
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
result[idst] = xv*post_reg[idst];
|
||||
}
|
||||
|
||||
device const char * rb = residual + i0*args.nb_r0 + it*args.nb_r2;
|
||||
FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) {
|
||||
const float rv = *(device const float *) (rb + isrc*args.nb_r1);
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
result[idst] = fma(rv, comb_reg[isrc][idst], result[idst]);
|
||||
}
|
||||
}
|
||||
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
*(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst];
|
||||
}
|
||||
}
|
||||
@@ -1,842 +0,0 @@
|
||||
#include "common.h"
|
||||
#include "dequantize.h"
|
||||
|
||||
constant bool FC_mul_mm_bc_inp [[function_constant(FC_MUL_MM + 0)]];
|
||||
constant bool FC_mul_mm_bc_out [[function_constant(FC_MUL_MM + 1)]];
|
||||
constant short FC_mul_mm_ne12 [[function_constant(FC_MUL_MM + 2)]];
|
||||
constant short FC_mul_mm_ne13 [[function_constant(FC_MUL_MM + 3)]];
|
||||
constant short FC_mul_mm_r2 [[function_constant(FC_MUL_MM + 4)]];
|
||||
constant short FC_mul_mm_r3 [[function_constant(FC_MUL_MM + 5)]];
|
||||
|
||||
// each block_q contains 16*nl weights
|
||||
#ifdef GGML_METAL_HAS_TENSOR
|
||||
template<
|
||||
typename SA, typename SA_4x4, typename SA_8x8,
|
||||
typename SB, typename SB_2x4, typename SB_8x8,
|
||||
typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread SA_4x4 &),
|
||||
typename T0, typename T0_4x4, typename T1, typename T1_2x4>
|
||||
kernel void kernel_mul_mm(
|
||||
constant ggml_metal_kargs_mul_mm & args,
|
||||
device const char * srcA,
|
||||
device const char * srcB,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig [[threadgroup_position_in_grid]],
|
||||
ushort tiitg [[thread_index_in_threadgroup]],
|
||||
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
|
||||
(void) sgitg;
|
||||
|
||||
// Matrix dimensions: A(M,K) x B(K,N) -> C(M,N)
|
||||
const int K = args.ne00;
|
||||
const int M = args.ne0;
|
||||
const int N = args.ne1;
|
||||
|
||||
// Batch dimension handling
|
||||
const int im = tgpig.z;
|
||||
const int i12 = im % FC_mul_mm_ne12;
|
||||
const int i13 = im / FC_mul_mm_ne12;
|
||||
|
||||
// Batch offsets for srcA and srcB
|
||||
const uint64_t offset0 = (i12/FC_mul_mm_r2)*args.nb02 + (i13/FC_mul_mm_r3)*args.nb03;
|
||||
|
||||
// Tile dimensions
|
||||
constexpr int NRB = SZ_SIMDGROUP * N_MM_BLOCK_X * N_MM_SIMD_GROUP_X;
|
||||
constexpr int NRA = SZ_SIMDGROUP * N_MM_BLOCK_Y * N_MM_SIMD_GROUP_Y;
|
||||
|
||||
// Tile offsets in output matrix
|
||||
const int ra = tgpig.y * NRA;
|
||||
const int rb = tgpig.x * NRB;
|
||||
|
||||
// Threadgroup memory for dequantized A tile only
|
||||
threadgroup SA * sa = (threadgroup SA *)(shmem);
|
||||
|
||||
// Work-item count for A loading
|
||||
constexpr int A_WORK_ITEMS = NRA * N_MM_NK;
|
||||
constexpr int NUM_THREADS = N_SIMDWIDTH * N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y;
|
||||
|
||||
// tA wraps threadgroup memory
|
||||
auto tA = tensor(sa, dextents<int32_t, 2>(N_MM_NK_TOTAL, NRA));
|
||||
|
||||
// tB wraps device memory directly
|
||||
device T1 * ptrB = (device T1 *)(srcB + args.nb12*i12 + args.nb13*i13);
|
||||
const int strideB = args.nb11 / sizeof(T1);
|
||||
auto tB = tensor(ptrB, dextents<int32_t, 2>(K, N), array<int, 2>({1, strideB}));
|
||||
|
||||
// Configure matmul operation
|
||||
mpp::tensor_ops::matmul2d<
|
||||
mpp::tensor_ops::matmul2d_descriptor(
|
||||
NRB, NRA, N_MM_NK_TOTAL, false, true, true,
|
||||
mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
|
||||
execution_simdgroups<N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y>> mm;
|
||||
|
||||
auto cT = mm.get_destination_cooperative_tensor<decltype(tB), decltype(tA), float>();
|
||||
|
||||
// Accumulate partial results over K dimension
|
||||
for (int loop_k = 0; loop_k < K; loop_k += N_MM_NK_TOTAL) {
|
||||
// === PHASE 1: Dequantization of A into threadgroup memory ===
|
||||
for (int work = tiitg; work < A_WORK_ITEMS; work += NUM_THREADS) {
|
||||
const int row = work / N_MM_NK;
|
||||
const int k_chunk = work % N_MM_NK;
|
||||
const int k_pos = loop_k + k_chunk * 16;
|
||||
const short k_base = k_chunk * 16;
|
||||
|
||||
// Bounds check: skip device read if row is out of matrix bounds
|
||||
if (ra + row < M) {
|
||||
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
|
||||
// Element-wise reads when K is not aligned (nb01 not aligned for half4x4/float4x4).
|
||||
// MSL spec Table 2.5: half4x4 requires 8-byte alignment. When K is odd,
|
||||
// nb01 = K*2 is not 8-byte aligned, so odd-row pointers are misaligned.
|
||||
// Mirrors the legacy kernel's existing guard.
|
||||
device const T0 * row_ptr = (device const T0 *)(srcA + args.nb01 * (ra + row) + offset0);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 16; i++) {
|
||||
sa[row * N_MM_NK_TOTAL + (k_base + i)] = (k_pos + i < K) ? (SA) row_ptr[k_pos + i] : (SA)0;
|
||||
}
|
||||
} else {
|
||||
const int block_idx = k_pos / (16 * nl);
|
||||
const short il = (k_pos / 16) % nl;
|
||||
|
||||
device const block_q * row_ptr = (device const block_q *)(srcA + args.nb01 * (ra + row) + offset0);
|
||||
|
||||
SA_4x4 temp_a;
|
||||
dequantize_func(row_ptr + block_idx, il, temp_a);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 16; i++) {
|
||||
// Zero-pad A for K positions beyond valid range (handles partial K iterations)
|
||||
sa[row * N_MM_NK_TOTAL + (k_base + i)] = (k_pos + i < K) ? temp_a[i/4][i%4] : (SA)0;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Zero-pad rows beyond matrix bounds
|
||||
FOR_UNROLL (short i = 0; i < 16; i++) {
|
||||
sa[row * N_MM_NK_TOTAL + (k_base + i)] = (SA)0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// === PHASE 2: Tensor matmul ===
|
||||
auto mA = tA.slice(0, 0);
|
||||
auto mB = tB.slice(loop_k, rb);
|
||||
|
||||
mm.run(mB, mA, cT);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
// Store result tile to output matrix (with batch offset)
|
||||
// cT.store handles bounds checking via tD's extents (M, N)
|
||||
device float * dstBatch = (device float *)dst + im * N * M;
|
||||
|
||||
auto tD = tensor(dstBatch, dextents<int32_t, 2>(M, N), array<int, 2>({1, M}));
|
||||
cT.store(tD.slice(ra, rb));
|
||||
}
|
||||
|
||||
#else
|
||||
|
||||
template<
|
||||
typename S0, typename S0_4x4, typename S0_8x8,
|
||||
typename S1, typename S1_2x4, typename S1_8x8,
|
||||
typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &),
|
||||
typename T0, typename T0_4x4, typename T1, typename T1_2x4>
|
||||
kernel void kernel_mul_mm(
|
||||
constant ggml_metal_kargs_mul_mm & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
|
||||
threadgroup S0 * sa = (threadgroup S0 *)(shmem);
|
||||
threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096);
|
||||
|
||||
constexpr int NR0 = 64;
|
||||
constexpr int NR1 = 32;
|
||||
|
||||
constexpr int NK = 32;
|
||||
constexpr int NL0 = NK/16;
|
||||
constexpr int NL1 = NK/8;
|
||||
|
||||
const int im = tgpig.z;
|
||||
const int r0 = tgpig.y*NR0;
|
||||
const int r1 = tgpig.x*NR1;
|
||||
|
||||
// if this block is of 64x32 shape or smaller
|
||||
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
|
||||
const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1;
|
||||
|
||||
// a thread shouldn't load data outside of the matrix
|
||||
const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63
|
||||
const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
|
||||
|
||||
const short il0 = (tiitg % NL0);
|
||||
|
||||
short il = il0;
|
||||
|
||||
const int i12 = im % FC_mul_mm_ne12;
|
||||
const int i13 = im / FC_mul_mm_ne12;
|
||||
|
||||
const uint64_t offset0 = (i12/FC_mul_mm_r2)*args.nb02 + (i13/FC_mul_mm_r3)*args.nb03;
|
||||
const short offset1 = il0/nl;
|
||||
|
||||
device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
|
||||
|
||||
const short iy = 8*(tiitg % NL1);
|
||||
|
||||
device const T1 * y = (device const T1 *)(src1
|
||||
+ args.nb13*i13
|
||||
+ args.nb12*i12
|
||||
+ args.nb11*(r1 + lr1)
|
||||
+ args.nb10*iy);
|
||||
|
||||
S0_8x8 ma[4];
|
||||
S1_8x8 mb[2];
|
||||
|
||||
simdgroup_float8x8 mc[8];
|
||||
|
||||
for (short i = 0; i < 8; i++){
|
||||
mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
||||
}
|
||||
|
||||
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
|
||||
// load data and store to threadgroup memory
|
||||
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// no need for dequantization
|
||||
for (short i = 0; i < 16; i++) {
|
||||
const short sx = 2*il0 + i/8;
|
||||
const short sy = (tiitg/NL0)/8;
|
||||
|
||||
//const short lx = i%8;
|
||||
//const short ly = (tiitg/NL0)%8;
|
||||
const short lx = (tiitg/NL0)%8;
|
||||
const short ly = i%8;
|
||||
|
||||
const short ib = 8*sx + sy;
|
||||
|
||||
*(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0;
|
||||
}
|
||||
} else {
|
||||
S0_4x4 temp_a;
|
||||
dequantize_func(x, il, temp_a);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 16; i++) {
|
||||
const short sx = 2*il0 + i/8;
|
||||
const short sy = (tiitg/NL0)/8;
|
||||
|
||||
//const short lx = i%8;
|
||||
//const short ly = (tiitg/NL0)%8;
|
||||
const short lx = (tiitg/NL0)%8;
|
||||
const short ly = i%8;
|
||||
|
||||
const short ib = 8*sx + sy;
|
||||
|
||||
// NOTE: this is massively slower.. WTF?
|
||||
//sa[64*ib + 8*ly + lx] = temp_a[i/4][i%4];
|
||||
|
||||
*(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_mul_mm_bc_inp) {
|
||||
for (short i = 0; i < 8; ++i) {
|
||||
const short sx = (tiitg%NL1);
|
||||
const short sy = (tiitg/NL1)/8;
|
||||
|
||||
const short lx = i;
|
||||
const short ly = (tiitg/NL1)%8;
|
||||
//const short lx = (tiitg/NL1)%8;
|
||||
//const short ly = i;
|
||||
|
||||
const short ib = 4*sx + sy;
|
||||
|
||||
*(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
|
||||
}
|
||||
} else {
|
||||
const short sx = (tiitg%NL1);
|
||||
const short sy = (tiitg/NL1)/8;
|
||||
|
||||
//const short dx = sx;
|
||||
//const short dy = sy;
|
||||
|
||||
const short ly = (tiitg/NL1)%8;
|
||||
|
||||
const short ib = 4*sx + sy;
|
||||
|
||||
*(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
|
||||
}
|
||||
|
||||
il = (il + 2 < nl) ? il + 2 : il % 2;
|
||||
x = (il < 2) ? x + (2 + nl - 1)/nl : x;
|
||||
|
||||
y += NK;
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// load matrices from threadgroup memory and conduct outer products
|
||||
threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
|
||||
threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
|
||||
|
||||
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
|
||||
simdgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 4; i++) {
|
||||
simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
|
||||
}
|
||||
|
||||
simdgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 2; i++) {
|
||||
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
|
||||
}
|
||||
|
||||
simdgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 8; i++){
|
||||
simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
|
||||
}
|
||||
|
||||
lsma += 8*64;
|
||||
lsmb += 4*64;
|
||||
}
|
||||
}
|
||||
|
||||
if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) {
|
||||
// if no bounds checks on the output are needed, we can directly write to device memory
|
||||
device float * C = (device float *) dst +
|
||||
(r0 + 32*(sgitg & 1)) + \
|
||||
(r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
|
||||
|
||||
for (short i = 0; i < 8; i++) {
|
||||
simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
|
||||
}
|
||||
} else {
|
||||
// block is smaller than 64x32, we should avoid writing data outside of the matrix
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
|
||||
|
||||
for (short i = 0; i < 8; i++) {
|
||||
simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (sgitg == 0) {
|
||||
for (int j = tiitg; j < nr1; j += NR1) {
|
||||
device float * D = (device float *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0;
|
||||
device float4 * D4 = (device float4 *) D;
|
||||
|
||||
threadgroup float * C = temp_str + (j*NR0);
|
||||
threadgroup float4 * C4 = (threadgroup float4 *) C;
|
||||
|
||||
int i = 0;
|
||||
for (; i < nr0/4; i++) {
|
||||
*(D4 + i) = *(C4 + i);
|
||||
}
|
||||
|
||||
i *= 4;
|
||||
for (; i < nr0; i++) {
|
||||
*(D + i) = *(C + i);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#endif // GGML_METAL_HAS_TENSOR
|
||||
|
||||
template<short ne20> // n_expert_used
|
||||
kernel void kernel_mul_mm_id_map0(
|
||||
constant ggml_metal_kargs_mul_mm_id_map0 & args,
|
||||
device const char * src2,
|
||||
device char * htpe,
|
||||
device char * hids,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
ushort tpitg[[thread_position_in_threadgroup]],
|
||||
ushort ntg[[threads_per_threadgroup]]) {
|
||||
const short ide = tpitg; // expert id
|
||||
|
||||
uint32_t n_all = 0;
|
||||
|
||||
device int32_t * ids_i32 = (device int32_t *) hids + ide*args.ne21;
|
||||
|
||||
for (int i21 = 0; i21 < args.ne21; i21 += ntg) { // n_tokens
|
||||
if (i21 + tpitg < args.ne21) {
|
||||
device const int32_t * src2_i32 = (device const int32_t *) (src2 + (i21 + tpitg)*args.nb21);
|
||||
|
||||
threadgroup uint16_t * sids = (threadgroup uint16_t *) shmem + tpitg*ne20;
|
||||
|
||||
#pragma unroll(ne20)
|
||||
for (short i20 = 0; i20 < ne20; i20++) {
|
||||
sids[i20] = src2_i32[i20];
|
||||
}
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (short t = 0; t < ntg; t++) {
|
||||
if (i21 + t >= args.ne21) {
|
||||
break;
|
||||
}
|
||||
|
||||
threadgroup const uint16_t * sids = (threadgroup const uint16_t *) shmem + t*ne20;
|
||||
|
||||
short sel = 0;
|
||||
#pragma unroll(ne20)
|
||||
for (short i20 = 0; i20 < ne20; i20++) {
|
||||
sel += (sids[i20] == ide)*(i20 + 1);
|
||||
}
|
||||
|
||||
ids_i32[n_all] = (i21 + t)*ne20 + sel - 1;
|
||||
|
||||
n_all += sel > 0;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
device uint32_t * tpe_u32 = (device uint32_t *) (htpe);
|
||||
tpe_u32[ide] = n_all;
|
||||
}
|
||||
|
||||
typedef decltype(kernel_mul_mm_id_map0<1>) kernel_mul_mm_id_map0_t;
|
||||
|
||||
template [[host_name("kernel_mul_mm_id_map0_ne20_1" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<1>;
|
||||
template [[host_name("kernel_mul_mm_id_map0_ne20_2" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<2>;
|
||||
template [[host_name("kernel_mul_mm_id_map0_ne20_4" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<4>;
|
||||
template [[host_name("kernel_mul_mm_id_map0_ne20_5" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<5>;
|
||||
template [[host_name("kernel_mul_mm_id_map0_ne20_6" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<6>;
|
||||
template [[host_name("kernel_mul_mm_id_map0_ne20_8" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<8>;
|
||||
template [[host_name("kernel_mul_mm_id_map0_ne20_10")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<10>;
|
||||
template [[host_name("kernel_mul_mm_id_map0_ne20_16")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<16>;
|
||||
template [[host_name("kernel_mul_mm_id_map0_ne20_22")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<22>;
|
||||
|
||||
template<typename S0, typename S0_4x4, typename S0_8x8, typename S1, typename S1_2x4, typename S1_8x8, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &), typename T0, typename T0_4x4, typename T1, typename T1_2x4>
|
||||
kernel void kernel_mul_mm_id(
|
||||
constant ggml_metal_kargs_mul_mm_id & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * htpe,
|
||||
device const char * hids,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
|
||||
threadgroup S0 * sa = (threadgroup S0 *)(shmem);
|
||||
threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096);
|
||||
|
||||
#ifdef GGML_METAL_HAS_TENSOR
|
||||
threadgroup float * sc = (threadgroup float *)(shmem);
|
||||
#endif
|
||||
|
||||
constexpr int NR0 = 64;
|
||||
constexpr int NR1 = 32;
|
||||
|
||||
constexpr int NK = 32;
|
||||
constexpr int NL0 = NK/16;
|
||||
constexpr int NL1 = NK/8;
|
||||
|
||||
const int im = tgpig.z; // expert
|
||||
const int r0 = tgpig.y*NR0;
|
||||
const int r1 = tgpig.x*NR1;
|
||||
|
||||
device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe);
|
||||
device const int32_t * ids_i32 = (device const int32_t *) (hids);
|
||||
|
||||
const int32_t neh1 = tpe_u32[im];
|
||||
|
||||
if (r1 >= neh1) {
|
||||
return;
|
||||
}
|
||||
|
||||
// if this block is of 64x32 shape or smaller
|
||||
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
|
||||
const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1;
|
||||
|
||||
// a thread shouldn't load data outside of the matrix
|
||||
const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63
|
||||
const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
|
||||
|
||||
const short il0 = (tiitg % NL0);
|
||||
|
||||
short il = il0;
|
||||
|
||||
const int id = ids_i32[im*args.ne21 + r1 + lr1];
|
||||
|
||||
const short i11 = (id % args.ne20) % args.ne11;
|
||||
const short i12 = (id / args.ne20);
|
||||
const short i13 = 0;
|
||||
|
||||
const uint64_t offset0 = im*args.nb02 + i13*args.nb03;
|
||||
const short offset1 = il0/nl;
|
||||
|
||||
device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
|
||||
|
||||
const short iy = 8*(tiitg % NL1);
|
||||
|
||||
device const T1 * y = (device const T1 *)(src1
|
||||
+ args.nb13*i13
|
||||
+ args.nb12*i12
|
||||
+ args.nb11*i11
|
||||
+ args.nb10*iy);
|
||||
|
||||
#ifndef GGML_METAL_HAS_TENSOR
|
||||
S0_8x8 ma[4];
|
||||
S1_8x8 mb[2];
|
||||
|
||||
simdgroup_float8x8 mc[8];
|
||||
|
||||
for (short i = 0; i < 8; i++){
|
||||
mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
|
||||
}
|
||||
#else
|
||||
auto tA = tensor<threadgroup S0, dextents<int32_t, 2>, tensor_inline>(sa, dextents<int32_t, 2>(NK, NR0));
|
||||
auto tB = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb, dextents<int32_t, 2>(NR1, NK ));
|
||||
|
||||
mpp::tensor_ops::matmul2d<
|
||||
mpp::tensor_ops::matmul2d_descriptor(NR1, NR0, NK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
|
||||
execution_simdgroups<4>> mm;
|
||||
|
||||
auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>();
|
||||
#endif
|
||||
|
||||
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
|
||||
#ifndef GGML_METAL_HAS_TENSOR
|
||||
// load data and store to threadgroup memory
|
||||
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// no need for dequantization
|
||||
for (short i = 0; i < 16; i++) {
|
||||
const short sx = 2*il0 + i/8;
|
||||
const short sy = (tiitg/NL0)/8;
|
||||
|
||||
//const short lx = i%8;
|
||||
//const short ly = (tiitg/NL0)%8;
|
||||
const short lx = (tiitg/NL0)%8;
|
||||
const short ly = i%8;
|
||||
|
||||
const short ib = 8*sx + sy;
|
||||
|
||||
*(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? (S0) *((device T0 *) x + i) : (S0) 0;
|
||||
}
|
||||
} else {
|
||||
S0_4x4 temp_a;
|
||||
dequantize_func(x, il, temp_a);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 16; i++) {
|
||||
const short sx = 2*il0 + i/8;
|
||||
const short sy = (tiitg/NL0)/8;
|
||||
|
||||
//const short lx = i%8;
|
||||
//const short ly = (tiitg/NL0)%8;
|
||||
const short lx = (tiitg/NL0)%8;
|
||||
const short ly = i%8;
|
||||
|
||||
const short ib = 8*sx + sy;
|
||||
|
||||
// NOTE: this is massively slower.. WTF?
|
||||
//sa[64*ib + 8*ly + lx] = temp_a[i/4][i%4];
|
||||
|
||||
*(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_mul_mm_bc_inp) {
|
||||
for (short i = 0; i < 8; ++i) {
|
||||
const short sx = (tiitg%NL1);
|
||||
const short sy = (tiitg/NL1)/8;
|
||||
|
||||
const short lx = i;
|
||||
const short ly = (tiitg/NL1)%8;
|
||||
//const short lx = (tiitg/NL1)%8;
|
||||
//const short ly = i;
|
||||
|
||||
const short ib = 4*sx + sy;
|
||||
|
||||
*(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
|
||||
}
|
||||
} else {
|
||||
const short sx = (tiitg%NL1);
|
||||
const short sy = (tiitg/NL1)/8;
|
||||
|
||||
//const short dx = sx;
|
||||
//const short dy = sy;
|
||||
|
||||
const short ly = (tiitg/NL1)%8;
|
||||
|
||||
const short ib = 4*sx + sy;
|
||||
|
||||
*(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
|
||||
}
|
||||
#else
|
||||
// load data and store to threadgroup memory
|
||||
if (is_same<T0_4x4, block_q>::value && FC_mul_mm_bc_inp) {
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// no need for dequantization
|
||||
for (short i = 0; i < 16; i++) {
|
||||
const short sx = 2*il0 + i/8;
|
||||
const short sy = (tiitg/NL0)/8;
|
||||
|
||||
const short lx = i%8;
|
||||
const short ly = (tiitg/NL0)%8;
|
||||
//const short lx = (tiitg/NL0)%8;
|
||||
//const short ly = i%8;
|
||||
|
||||
*(sa + NK*(8*sy + ly) + 8*sx + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0;
|
||||
}
|
||||
} else {
|
||||
S0_4x4 temp_a;
|
||||
dequantize_func(x, il, temp_a);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 16; i++) {
|
||||
const short sx = 2*il0 + i/8;
|
||||
const short sy = (tiitg/NL0)/8;
|
||||
|
||||
const short lx = i%8;
|
||||
const short ly = (tiitg/NL0)%8;
|
||||
//const short lx = (tiitg/NL0)%8;
|
||||
//const short ly = i%8;
|
||||
|
||||
*(sa + NK*(8*sy + ly) + 8*sx + lx) = temp_a[i/4][i%4];
|
||||
}
|
||||
}
|
||||
|
||||
if (FC_mul_mm_bc_inp) {
|
||||
for (short i = 0; i < 8; ++i) {
|
||||
const short sx = (tiitg%NL1);
|
||||
const short sy = (tiitg/NL1)/8;
|
||||
|
||||
const short lx = i;
|
||||
const short ly = (tiitg/NL1)%8;
|
||||
//const short lx = (tiitg/NL1)%8;
|
||||
//const short ly = i;
|
||||
|
||||
*(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
|
||||
}
|
||||
} else {
|
||||
const short sx = (tiitg%NL1);
|
||||
const short sy = (tiitg/NL1)/8;
|
||||
|
||||
//const short lx = i;
|
||||
const short ly = (tiitg/NL1)%8;
|
||||
//const short lx = (tiitg/NL1)%8;
|
||||
//const short ly = i;
|
||||
|
||||
*(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)(*((device T1_2x4 *) y));
|
||||
}
|
||||
#endif
|
||||
|
||||
il = (il + 2 < nl) ? il + 2 : il % 2;
|
||||
x = (il < 2) ? x + (2 + nl - 1)/nl : x;
|
||||
|
||||
y += NK;
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
#ifndef GGML_METAL_HAS_TENSOR
|
||||
// load matrices from threadgroup memory and conduct outer products
|
||||
threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
|
||||
threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
|
||||
|
||||
FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
|
||||
simdgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 4; i++) {
|
||||
simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
|
||||
}
|
||||
|
||||
simdgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 2; i++) {
|
||||
simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
|
||||
}
|
||||
|
||||
simdgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
FOR_UNROLL (short i = 0; i < 8; i++){
|
||||
simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
|
||||
}
|
||||
|
||||
lsma += 8*64;
|
||||
lsmb += 4*64;
|
||||
}
|
||||
#else
|
||||
auto sA = tA.slice(0, 0);
|
||||
auto sB = tB.slice(0, 0);
|
||||
|
||||
mm.run(sB, sA, cT);
|
||||
#endif
|
||||
}
|
||||
|
||||
// block is smaller than 64x32, we should avoid writing data outside of the matrix
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
#ifdef GGML_METAL_HAS_TENSOR
|
||||
auto tC = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc, dextents<int32_t, 2>(NR0, NR1));
|
||||
cT.store(tC);
|
||||
#else
|
||||
threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
|
||||
|
||||
for (short i = 0; i < 8; i++) {
|
||||
simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
|
||||
}
|
||||
#endif
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (short j = sgitg; j < nr1; j += 4) {
|
||||
const int id = ids_i32[im*args.ne21 + r1 + j];
|
||||
|
||||
const short ide = id % args.ne20;
|
||||
const short idt = id / args.ne20;
|
||||
|
||||
device float * D = (device float *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0;
|
||||
device float4 * D4 = (device float4 *) D;
|
||||
|
||||
threadgroup float * C = (threadgroup float *) shmem + j*NR0;
|
||||
threadgroup float4 * C4 = (threadgroup float4 *) C;
|
||||
|
||||
int i = tiisg;
|
||||
for (; i < nr0/4; i += 32) {
|
||||
*(D4 + i) = *(C4 + i);
|
||||
}
|
||||
|
||||
i = (4*(nr0/4)) + tiisg;
|
||||
for (; i < nr0; i += 32) {
|
||||
*(D + i) = *(C + i);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// matrix-matrix multiplication
|
||||
//
|
||||
|
||||
typedef decltype(kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>) mul_mm_t;
|
||||
|
||||
template [[host_name("kernel_mul_mm_f32_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_f16_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, half4x4, 1, dequantize_f16, half, half4x4, float, float2x4>;
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_mul_mm_bf16_f32")]] kernel mul_mm_t kernel_mul_mm<bfloat, bfloat4x4, simdgroup_bfloat8x8, bfloat, bfloat2x4, simdgroup_bfloat8x8, bfloat4x4, 1, dequantize_bf16, bfloat, bfloat4x4, float, float2x4>;
|
||||
#endif
|
||||
template [[host_name("kernel_mul_mm_q1_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q1_0, 8, dequantize_q1_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q2_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_0, 4, dequantize_q2_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q4_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_0, 2, dequantize_q4_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q4_1_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_1, 2, dequantize_q4_1, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q5_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_0, 2, dequantize_q5_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q5_1_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_1, 2, dequantize_q5_1, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q8_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_mxfp4_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q2_K_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q3_K_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q3_K, QK_NL, dequantize_q3_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q4_K_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q5_K_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_q6_K_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq2_xxs_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq2_xs_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xs, QK_NL, dequantize_iq2_xs, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq3_xxs_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq3_s_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq3_s, QK_NL, dequantize_iq3_s, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq2_s_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_s, QK_NL, dequantize_iq2_s, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq1_s_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_s, QK_NL, dequantize_iq1_s, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq1_m_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_m, QK_NL, dequantize_iq1_m, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq4_nl_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_nl, 2, dequantize_iq4_nl, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq4_xs_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_xs, QK_NL, dequantize_iq4_xs, float, float4x4, float, float2x4>;
|
||||
|
||||
template [[host_name("kernel_mul_mm_f32_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_f16_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, half4x4, 1, dequantize_f16, half, half4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q1_0_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q1_0, 8, dequantize_q1_0, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q2_0_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_0, 4, dequantize_q2_0, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q4_0_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_0, 2, dequantize_q4_0, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q4_1_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_1, 2, dequantize_q4_1, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q5_0_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_0, 2, dequantize_q5_0, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q5_1_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_1, 2, dequantize_q5_1, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q8_0_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_mxfp4_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q2_K_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q3_K_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q3_K, QK_NL, dequantize_q3_K, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q4_K_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q5_K_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_q6_K_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq2_xxs_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq2_xs_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xs, QK_NL, dequantize_iq2_xs, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq3_xxs_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq3_s_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq3_s, QK_NL, dequantize_iq3_s, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq2_s_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_s, QK_NL, dequantize_iq2_s, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq1_s_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_s, QK_NL, dequantize_iq1_s, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq1_m_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_m, QK_NL, dequantize_iq1_m, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq4_nl_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_nl, 2, dequantize_iq4_nl, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_iq4_xs_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_xs, QK_NL, dequantize_iq4_xs, float, float4x4, half, half2x4>;
|
||||
|
||||
//
|
||||
// indirect matrix-matrix multiplication
|
||||
//
|
||||
|
||||
typedef decltype(kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>) mul_mm_id;
|
||||
|
||||
template [[host_name("kernel_mul_mm_id_f32_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_f16_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, half4x4, 1, dequantize_f16, half, half4x4, float, float2x4>;
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_mul_mm_id_bf16_f32")]] kernel mul_mm_id kernel_mul_mm_id<bfloat, bfloat4x4, simdgroup_bfloat8x8, bfloat, bfloat2x4, simdgroup_bfloat8x8, bfloat4x4, 1, dequantize_bf16, bfloat, bfloat4x4, float, float2x4>;
|
||||
#endif
|
||||
template [[host_name("kernel_mul_mm_id_q1_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q1_0, 8, dequantize_q1_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q2_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_0, 4, dequantize_q2_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q4_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_0, 2, dequantize_q4_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q4_1_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_1, 2, dequantize_q4_1, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q5_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_0, 2, dequantize_q5_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q5_1_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_1, 2, dequantize_q5_1, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q8_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_mxfp4_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q2_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q3_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q3_K, QK_NL, dequantize_q3_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q4_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q5_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q6_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq2_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq2_xs_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xs, QK_NL, dequantize_iq2_xs, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq3_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq3_s_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq3_s, QK_NL, dequantize_iq3_s, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq2_s_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_s, QK_NL, dequantize_iq2_s, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq1_s_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_s, QK_NL, dequantize_iq1_s, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq1_m_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_m, QK_NL, dequantize_iq1_m, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq4_nl_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_nl, 2, dequantize_iq4_nl, float, float4x4, float, float2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq4_xs_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_xs, QK_NL, dequantize_iq4_xs, float, float4x4, float, float2x4>;
|
||||
|
||||
template [[host_name("kernel_mul_mm_id_f32_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_f16_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, half4x4, 1, dequantize_f16, half, half4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q1_0_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q1_0, 8, dequantize_q1_0, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q2_0_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_0, 4, dequantize_q2_0, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q4_0_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_0, 2, dequantize_q4_0, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q4_1_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_1, 2, dequantize_q4_1, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q5_0_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_0, 2, dequantize_q5_0, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q5_1_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_1, 2, dequantize_q5_1, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q8_0_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_mxfp4_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_mxfp4, 2, dequantize_mxfp4, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q2_K_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q3_K_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q3_K, QK_NL, dequantize_q3_K, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q4_K_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q5_K_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_q6_K_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq2_xxs_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq2_xs_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xs, QK_NL, dequantize_iq2_xs, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq3_xxs_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq3_xxs, QK_NL, dequantize_iq3_xxs, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq3_s_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq3_s, QK_NL, dequantize_iq3_s, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq2_s_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_s, QK_NL, dequantize_iq2_s, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq1_s_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_s, QK_NL, dequantize_iq1_s, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq1_m_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_m, QK_NL, dequantize_iq1_m, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq4_nl_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_nl, 2, dequantize_iq4_nl, float, float4x4, half, half2x4>;
|
||||
template [[host_name("kernel_mul_mm_id_iq4_xs_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_xs, QK_NL, dequantize_iq4_xs, float, float4x4, half, half2x4>;
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,308 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
// F == 1 : norm (no fuse)
|
||||
// F == 2 : norm + mul
|
||||
// F == 3 : norm + mul + add
|
||||
template <typename T, short F>
|
||||
kernel void kernel_norm_fuse_impl(
|
||||
constant ggml_metal_kargs_norm & args,
|
||||
device const char * src0,
|
||||
device const char * src1_0,
|
||||
device const char * src1_1,
|
||||
device char * dst,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
const int i01 = tgpig.x;
|
||||
const int i02 = tgpig.y;
|
||||
const int i03 = tgpig.z;
|
||||
|
||||
device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
|
||||
|
||||
device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
|
||||
device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);
|
||||
|
||||
T sumft(0.0f);
|
||||
|
||||
float sumf = 0.0f;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
sumft += x[i00];
|
||||
}
|
||||
sumf = dot(sumft, T(1.0f));
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_f32[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
const float mean = sumf/args.ne00;
|
||||
|
||||
device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
|
||||
|
||||
sumf = 0.0f;
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
y[i00] = x[i00] - mean;
|
||||
sumf += dot(y[i00], y[i00]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_f32[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
const float variance = sumf/args.ne00;
|
||||
|
||||
const float scale = 1.0f/sqrt(variance + args.eps);
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
if (F == 1) {
|
||||
y[i00] = (y[i00]*scale);
|
||||
}
|
||||
if (F == 2) {
|
||||
y[i00] = (y[i00]*scale)*f0[i00];
|
||||
}
|
||||
if (F == 3) {
|
||||
y[i00] = (y[i00]*scale)*f0[i00] + f1[i00];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_norm_fuse_impl<float4, 1>) kernel_norm_fuse_t;
|
||||
|
||||
template [[host_name("kernel_norm_f32")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 1>;
|
||||
template [[host_name("kernel_norm_mul_f32")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 2>;
|
||||
template [[host_name("kernel_norm_mul_add_f32")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 3>;
|
||||
|
||||
template [[host_name("kernel_norm_f32_4")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 1>;
|
||||
template [[host_name("kernel_norm_mul_f32_4")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 2>;
|
||||
template [[host_name("kernel_norm_mul_add_f32_4")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 3>;
|
||||
|
||||
// F == 1 : rms_norm (no fuse)
|
||||
// F == 2 : rms_norm + mul
|
||||
// F == 3 : rms_norm + mul + add
|
||||
template <typename T, short F>
|
||||
kernel void kernel_rms_norm_fuse_impl(
|
||||
constant ggml_metal_kargs_norm & args,
|
||||
device const char * src0,
|
||||
device const char * src1_0,
|
||||
device const char * src1_1,
|
||||
device char * dst,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
const int i01 = tgpig.x;
|
||||
const int i02 = tgpig.y;
|
||||
const int i03 = tgpig.z;
|
||||
|
||||
device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
|
||||
|
||||
device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
|
||||
device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);
|
||||
|
||||
float sumf = 0.0f;
|
||||
|
||||
// parallel sum
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
sumf += dot(x[i00], x[i00]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_f32[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
const float mean = sumf/args.ne00;
|
||||
const float scale = 1.0f/sqrt(mean + args.eps);
|
||||
|
||||
device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
|
||||
for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
|
||||
if (F == 1) {
|
||||
y[i00] = (x[i00]*scale);
|
||||
}
|
||||
if (F == 2) {
|
||||
y[i00] = (x[i00]*scale)*f0[i00];
|
||||
}
|
||||
if (F == 3) {
|
||||
y[i00] = (x[i00]*scale)*f0[i00] + f1[i00];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_rms_norm_fuse_impl<float4, 1>) kernel_rms_norm_fuse_t;
|
||||
|
||||
template [[host_name("kernel_rms_norm_f32")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 1>;
|
||||
template [[host_name("kernel_rms_norm_mul_f32")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 2>;
|
||||
template [[host_name("kernel_rms_norm_mul_add_f32")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 3>;
|
||||
|
||||
template [[host_name("kernel_rms_norm_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 1>;
|
||||
template [[host_name("kernel_rms_norm_mul_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 2>;
|
||||
template [[host_name("kernel_rms_norm_mul_add_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 3>;
|
||||
|
||||
template <typename T0, typename T>
|
||||
kernel void kernel_l2_norm_impl(
|
||||
constant ggml_metal_kargs_l2_norm & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i03 = tgpig.z;
|
||||
const int i02 = tgpig.y;
|
||||
const int i01 = tgpig.x;
|
||||
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
device const T0 * x = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
|
||||
device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
|
||||
|
||||
float sumf = 0.0f;
|
||||
|
||||
// parallel sum
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += ntg.x) {
|
||||
sumf += dot(x[i00], x[i00]);
|
||||
}
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_f32[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
const float scale = 1.0f/max(sqrt(sumf), args.eps);
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += ntg.x) {
|
||||
y[i00] = x[i00] * scale;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_l2_norm_impl<float, float>) kernel_l2_norm_t;
|
||||
|
||||
template [[host_name("kernel_l2_norm_f32_f32")]] kernel kernel_l2_norm_t kernel_l2_norm_impl<float, float>;
|
||||
template [[host_name("kernel_l2_norm_f32_f32_4")]] kernel kernel_l2_norm_t kernel_l2_norm_impl<float4, float4>;
|
||||
|
||||
kernel void kernel_group_norm_f32(
|
||||
constant ggml_metal_kargs_group_norm & args,
|
||||
device const float * src0,
|
||||
device float * dst,
|
||||
threadgroup float * buf [[threadgroup(0)]],
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
uint tpitg[[thread_position_in_threadgroup]],
|
||||
uint sgitg[[simdgroup_index_in_threadgroup]],
|
||||
uint tiisg[[thread_index_in_simdgroup]],
|
||||
uint ntg[[threads_per_threadgroup]]) {
|
||||
const int64_t ne = args.ne00*args.ne01*args.ne02;
|
||||
const int64_t gs = args.ne00*args.ne01*((args.ne02 + args.ngrp - 1) / args.ngrp);
|
||||
|
||||
int start = tgpig * gs;
|
||||
int end = start + gs;
|
||||
|
||||
start += tpitg;
|
||||
|
||||
if (end >= ne) {
|
||||
end = ne;
|
||||
}
|
||||
|
||||
float tmp = 0.0f; // partial sum for thread in warp
|
||||
|
||||
for (int j = start; j < end; j += ntg) {
|
||||
tmp += src0[j];
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
tmp = simd_sum(tmp);
|
||||
if (ntg > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = tmp;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
tmp = buf[tiisg];
|
||||
tmp = simd_sum(tmp);
|
||||
}
|
||||
|
||||
const float mean = tmp / gs;
|
||||
tmp = 0.0f;
|
||||
|
||||
for (int j = start; j < end; j += ntg) {
|
||||
float xi = src0[j] - mean;
|
||||
dst[j] = xi;
|
||||
tmp += xi * xi;
|
||||
}
|
||||
|
||||
tmp = simd_sum(tmp);
|
||||
if (ntg > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = tmp;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
tmp = buf[tiisg];
|
||||
tmp = simd_sum(tmp);
|
||||
}
|
||||
|
||||
const float variance = tmp / gs;
|
||||
const float scale = 1.0f/sqrt(variance + args.eps);
|
||||
for (int j = start; j < end; j += ntg) {
|
||||
dst[j] *= scale;
|
||||
}
|
||||
}
|
||||
@@ -1,148 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
kernel void kernel_pool_2d_max_f32(
|
||||
constant ggml_metal_kargs_pool_2d & args,
|
||||
device const float * src0,
|
||||
device float * dst,
|
||||
uint gid[[thread_position_in_grid]]) {
|
||||
|
||||
if (gid >= args.np) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int idx = gid;
|
||||
const int I_HW = args.IH * args.IW;
|
||||
const int O_HW = args.OH * args.OW;
|
||||
const int nc = idx / O_HW;
|
||||
const int cur_oh = idx % O_HW / args.OW;
|
||||
const int cur_ow = idx % O_HW % args.OW;
|
||||
|
||||
device const float * i_ptr = src0 + nc * I_HW;
|
||||
device float * o_ptr = dst + nc * O_HW;
|
||||
|
||||
const int start_h = cur_oh * args.s1 - args.p1;
|
||||
const int bh = MAX(0, start_h);
|
||||
const int eh = MIN(args.IH, start_h + args.k1);
|
||||
const int start_w = cur_ow * args.s0 - args.p0;
|
||||
const int bw = MAX(0, start_w);
|
||||
const int ew = MIN(args.IW, start_w + args.k0);
|
||||
|
||||
float res = -INFINITY;
|
||||
|
||||
for (int i = bh; i < eh; i += 1) {
|
||||
for (int j = bw; j < ew; j += 1) {
|
||||
res = MAX(res, i_ptr[i * args.IW + j]);
|
||||
}
|
||||
}
|
||||
|
||||
o_ptr[cur_oh * args.OW + cur_ow] = res;
|
||||
}
|
||||
|
||||
kernel void kernel_pool_2d_avg_f32(
|
||||
constant ggml_metal_kargs_pool_2d & args,
|
||||
device const float * src0,
|
||||
device float * dst,
|
||||
uint gid[[thread_position_in_grid]]) {
|
||||
|
||||
if (gid >= args.np) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int idx = gid;
|
||||
const int I_HW = args.IH * args.IW;
|
||||
const int O_HW = args.OH * args.OW;
|
||||
const int nc = idx / O_HW;
|
||||
const int cur_oh = idx % O_HW / args.OW;
|
||||
const int cur_ow = idx % O_HW % args.OW;
|
||||
|
||||
device const float * i_ptr = src0 + nc * I_HW;
|
||||
device float * o_ptr = dst + nc * O_HW;
|
||||
|
||||
const int start_h = cur_oh * args.s1 - args.p1;
|
||||
const int bh = MAX(0, start_h);
|
||||
const int eh = MIN(args.IH, start_h + args.k1);
|
||||
const int start_w = cur_ow * args.s0 - args.p0;
|
||||
const int bw = MAX(0, start_w);
|
||||
const int ew = MIN(args.IW, start_w + args.k0);
|
||||
// const float scale = 1. / ((eh - bh) * (ew - bw));
|
||||
const float scale = 1. / (args.k0 * args.k1);
|
||||
|
||||
float res = 0;
|
||||
|
||||
for (int i = bh; i < eh; i += 1) {
|
||||
for (int j = bw; j < ew; j += 1) {
|
||||
float cur = i_ptr[i * args.IW + j];
|
||||
res += cur * scale;
|
||||
}
|
||||
}
|
||||
|
||||
o_ptr[cur_oh * args.OW + cur_ow] = res;
|
||||
}
|
||||
|
||||
|
||||
kernel void kernel_pool_1d_max_f32(
|
||||
constant ggml_metal_kargs_pool_1d & args,
|
||||
device const float * src,
|
||||
device float * dst,
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
|
||||
if (gid >= args.np) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int ow = (int)gid % args.OW;
|
||||
const int row = (int)gid / args.OW;
|
||||
|
||||
const int base = ow * args.s0 - args.p0;
|
||||
|
||||
float acc = -INFINITY;
|
||||
|
||||
const int src_off = row * args.IW;
|
||||
const int dst_off = row * args.OW;
|
||||
|
||||
for (int ki = 0; ki < args.k0; ++ki) {
|
||||
int j = base + ki;
|
||||
if (j < 0 || j >= args.IW){
|
||||
continue;
|
||||
}
|
||||
float v = src[src_off + j];
|
||||
acc = max(acc, v);
|
||||
}
|
||||
|
||||
dst[dst_off + ow] = acc;
|
||||
}
|
||||
|
||||
kernel void kernel_pool_1d_avg_f32(
|
||||
constant ggml_metal_kargs_pool_1d & args,
|
||||
device const float * src,
|
||||
device float * dst,
|
||||
uint gid [[thread_position_in_grid]]
|
||||
) {
|
||||
|
||||
if (gid >= args.np) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int ow = (int)gid % args.OW;
|
||||
const int row = (int)gid / args.OW;
|
||||
|
||||
const int base = ow * args.s0 - args.p0;
|
||||
|
||||
float acc = 0.0f;
|
||||
int cnt = 0;
|
||||
|
||||
const int src_off = row * args.IW;
|
||||
const int dst_off = row * args.OW;
|
||||
|
||||
for (int ki = 0; ki < args.k0; ++ki) {
|
||||
const int j = base + ki;
|
||||
if (j < 0 || j >= args.IW) {
|
||||
continue;
|
||||
}
|
||||
acc += src[src_off + j];
|
||||
cnt += 1;
|
||||
}
|
||||
|
||||
dst[dst_off + ow] = (cnt > 0) ? (acc / (float)cnt) : 0.0f;
|
||||
}
|
||||
@@ -1,234 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
#include "common.h"
|
||||
|
||||
void quantize_q1_0(device const float * src, device block_q1_0 & dst) {
|
||||
float sum_abs = 0.0f;
|
||||
for (int j = 0; j < QK1_0; j++) {
|
||||
sum_abs += fabs(src[j]);
|
||||
}
|
||||
dst.d = sum_abs / QK1_0;
|
||||
|
||||
for (int j = 0; j < QK1_0 / 8; j++) {
|
||||
dst.qs[j] = 0;
|
||||
}
|
||||
for (int j = 0; j < QK1_0; j++) {
|
||||
if (src[j] >= 0.0f) {
|
||||
dst.qs[j / 8] |= (1 << (j % 8));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void quantize_q2_0(device const float * src, device block_q2_0 & dst) {
|
||||
float amax = 0.0f;
|
||||
for (int j = 0; j < QK2_0; j++) {
|
||||
float a = fabs(src[j]);
|
||||
if (a > amax) amax = a;
|
||||
}
|
||||
const float d = amax;
|
||||
dst.d = d;
|
||||
|
||||
const float id = d > 0.0f ? 1.0f / d : 0.0f;
|
||||
|
||||
for (int j = 0; j < QK2_0 / 4; j++) {
|
||||
dst.qs[j] = 0;
|
||||
}
|
||||
for (int j = 0; j < QK2_0; j++) {
|
||||
int q = (int)round(src[j] * id) + 1;
|
||||
q = max(0, min(3, q));
|
||||
dst.qs[j / 4] |= (q << (2 * (j % 4)));
|
||||
}
|
||||
}
|
||||
|
||||
void quantize_q4_0(device const float * src, device block_q4_0 & dst) {
|
||||
#pragma METAL fp math_mode(safe)
|
||||
float amax = 0.0f; // absolute max
|
||||
float max = 0.0f;
|
||||
|
||||
for (int j = 0; j < QK4_0; j++) {
|
||||
const float v = src[j];
|
||||
if (amax < fabs(v)) {
|
||||
amax = fabs(v);
|
||||
max = v;
|
||||
}
|
||||
}
|
||||
|
||||
const float d = max / -8;
|
||||
const float id = d ? 1.0f/d : 0.0f;
|
||||
|
||||
dst.d = d;
|
||||
|
||||
for (int j = 0; j < QK4_0/2; ++j) {
|
||||
const float x0 = src[0 + j]*id;
|
||||
const float x1 = src[QK4_0/2 + j]*id;
|
||||
|
||||
const uint8_t xi0 = MIN(15, (int8_t)(x0 + 8.5f));
|
||||
const uint8_t xi1 = MIN(15, (int8_t)(x1 + 8.5f));
|
||||
|
||||
dst.qs[j] = xi0;
|
||||
dst.qs[j] |= xi1 << 4;
|
||||
}
|
||||
}
|
||||
|
||||
void quantize_q4_1(device const float * src, device block_q4_1 & dst) {
|
||||
#pragma METAL fp math_mode(safe)
|
||||
float min = FLT_MAX;
|
||||
float max = -FLT_MAX;
|
||||
|
||||
for (int j = 0; j < QK4_1; j++) {
|
||||
const float v = src[j];
|
||||
if (min > v) min = v;
|
||||
if (max < v) max = v;
|
||||
}
|
||||
|
||||
const float d = (max - min) / ((1 << 4) - 1);
|
||||
const float id = d ? 1.0f/d : 0.0f;
|
||||
|
||||
dst.d = d;
|
||||
dst.m = min;
|
||||
|
||||
for (int j = 0; j < QK4_1/2; ++j) {
|
||||
const float x0 = (src[0 + j] - min)*id;
|
||||
const float x1 = (src[QK4_1/2 + j] - min)*id;
|
||||
|
||||
const uint8_t xi0 = MIN(15, (int8_t)(x0 + 0.5f));
|
||||
const uint8_t xi1 = MIN(15, (int8_t)(x1 + 0.5f));
|
||||
|
||||
dst.qs[j] = xi0;
|
||||
dst.qs[j] |= xi1 << 4;
|
||||
}
|
||||
}
|
||||
|
||||
void quantize_q5_0(device const float * src, device block_q5_0 & dst) {
|
||||
#pragma METAL fp math_mode(safe)
|
||||
float amax = 0.0f; // absolute max
|
||||
float max = 0.0f;
|
||||
|
||||
for (int j = 0; j < QK5_0; j++) {
|
||||
const float v = src[j];
|
||||
if (amax < fabs(v)) {
|
||||
amax = fabs(v);
|
||||
max = v;
|
||||
}
|
||||
}
|
||||
|
||||
const float d = max / -16;
|
||||
const float id = d ? 1.0f/d : 0.0f;
|
||||
|
||||
dst.d = d;
|
||||
|
||||
uint32_t qh = 0;
|
||||
for (int j = 0; j < QK5_0/2; ++j) {
|
||||
const float x0 = src[0 + j]*id;
|
||||
const float x1 = src[QK5_0/2 + j]*id;
|
||||
|
||||
const uint8_t xi0 = MIN(31, (int8_t)(x0 + 16.5f));
|
||||
const uint8_t xi1 = MIN(31, (int8_t)(x1 + 16.5f));
|
||||
|
||||
dst.qs[j] = (xi0 & 0xf) | ((xi1 & 0xf) << 4);
|
||||
qh |= ((xi0 & 0x10u) >> 4) << (j + 0);
|
||||
qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2);
|
||||
}
|
||||
|
||||
thread const uint8_t * qh8 = (thread const uint8_t *)&qh;
|
||||
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
dst.qh[j] = qh8[j];
|
||||
}
|
||||
}
|
||||
|
||||
void quantize_q5_1(device const float * src, device block_q5_1 & dst) {
|
||||
#pragma METAL fp math_mode(safe)
|
||||
float max = src[0];
|
||||
float min = src[0];
|
||||
|
||||
for (int j = 1; j < QK5_1; j++) {
|
||||
const float v = src[j];
|
||||
min = v < min ? v : min;
|
||||
max = v > max ? v : max;
|
||||
}
|
||||
|
||||
const float d = (max - min) / 31;
|
||||
const float id = d ? 1.0f/d : 0.0f;
|
||||
|
||||
dst.d = d;
|
||||
dst.m = min;
|
||||
|
||||
uint32_t qh = 0;
|
||||
for (int j = 0; j < QK5_1/2; ++j) {
|
||||
const float x0 = (src[0 + j] - min)*id;
|
||||
const float x1 = (src[QK5_1/2 + j] - min)*id;
|
||||
|
||||
const uint8_t xi0 = (uint8_t)(x0 + 0.5f);
|
||||
const uint8_t xi1 = (uint8_t)(x1 + 0.5f);
|
||||
|
||||
dst.qs[j] = (xi0 & 0xf) | ((xi1 & 0xf) << 4);
|
||||
qh |= ((xi0 & 0x10u) >> 4) << (j + 0);
|
||||
qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_1/2);
|
||||
}
|
||||
|
||||
thread const uint8_t * qh8 = (thread const uint8_t *)&qh;
|
||||
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
dst.qh[j] = qh8[j];
|
||||
}
|
||||
}
|
||||
|
||||
void quantize_q8_0(device const float * src, device block_q8_0 & dst) {
|
||||
#pragma METAL fp math_mode(safe)
|
||||
float amax = 0.0f; // absolute max
|
||||
|
||||
for (int j = 0; j < QK8_0; j++) {
|
||||
const float v = src[j];
|
||||
amax = MAX(amax, fabs(v));
|
||||
}
|
||||
|
||||
const float d = amax / ((1 << 7) - 1);
|
||||
const float id = d ? 1.0f/d : 0.0f;
|
||||
|
||||
dst.d = d;
|
||||
|
||||
for (int j = 0; j < QK8_0; ++j) {
|
||||
const float x0 = src[j]*id;
|
||||
|
||||
dst.qs[j] = round(x0);
|
||||
}
|
||||
}
|
||||
|
||||
void quantize_iq4_nl(device const float * src, device block_iq4_nl & dst) {
|
||||
#pragma METAL fp math_mode(safe)
|
||||
float amax = 0.0f; // absolute max
|
||||
float max = 0.0f;
|
||||
|
||||
for (int j = 0; j < QK4_NL; j++) {
|
||||
const float v = src[j];
|
||||
if (amax < fabs(v)) {
|
||||
amax = fabs(v);
|
||||
max = v;
|
||||
}
|
||||
}
|
||||
|
||||
const float d = max / kvalues_iq4nl_f[0];
|
||||
const float id = d ? 1.0f/d : 0.0f;
|
||||
|
||||
float sumqx = 0, sumq2 = 0;
|
||||
for (int j = 0; j < QK4_NL/2; ++j) {
|
||||
const float x0 = src[0 + j]*id;
|
||||
const float x1 = src[QK4_NL/2 + j]*id;
|
||||
|
||||
const uint8_t xi0 = best_index_int8(16, kvalues_iq4nl_f, x0);
|
||||
const uint8_t xi1 = best_index_int8(16, kvalues_iq4nl_f, x1);
|
||||
|
||||
dst.qs[j] = xi0 | (xi1 << 4);
|
||||
|
||||
const float v0 = kvalues_iq4nl_f[xi0];
|
||||
const float v1 = kvalues_iq4nl_f[xi1];
|
||||
const float w0 = src[0 + j]*src[0 + j];
|
||||
const float w1 = src[QK4_NL/2 + j]*src[QK4_NL/2 + j];
|
||||
sumqx += w0*v0*src[j] + w1*v1*src[QK4_NL/2 + j];
|
||||
sumq2 += w0*v0*v0 + w1*v1*v1;
|
||||
|
||||
}
|
||||
|
||||
dst.d = sumq2 > 0 ? sumqx/sumq2 : d;
|
||||
}
|
||||
@@ -1,393 +0,0 @@
|
||||
#include "common.h"
|
||||
#include "dequantize.h"
|
||||
#include "quantize.h"
|
||||
|
||||
template<typename T0, typename T1>
|
||||
kernel void kernel_cpy_t_t(
|
||||
constant ggml_metal_kargs_cpy & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig[2];
|
||||
const int32_t i02 = tgpig[1];
|
||||
const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
|
||||
const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
|
||||
|
||||
if (i01 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
|
||||
|
||||
const int32_t i3 = n/(args.ne2*args.ne1*args.ne0);
|
||||
const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0);
|
||||
const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
|
||||
const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
|
||||
|
||||
device T1 * dst_data = (device T1 *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.ne00;) {
|
||||
device const T0 * src = (device T0 *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
|
||||
dst_data[i00] = (T1) src[0];
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_cpy_t_t<float, float>) kernel_cpy_t;
|
||||
|
||||
template [[host_name("kernel_cpy_f32_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<float, float>;
|
||||
template [[host_name("kernel_cpy_f32_f16")]] kernel kernel_cpy_t kernel_cpy_t_t<float, half>;
|
||||
template [[host_name("kernel_cpy_f32_i32")]] kernel kernel_cpy_t kernel_cpy_t_t<float, int32_t>;
|
||||
template [[host_name("kernel_cpy_i32_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<int32_t, float>;
|
||||
template [[host_name("kernel_cpy_i32_i32")]] kernel kernel_cpy_t kernel_cpy_t_t<int32_t, int32_t>;
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_cpy_f32_bf16")]] kernel kernel_cpy_t kernel_cpy_t_t<float, bfloat>;
|
||||
#endif
|
||||
template [[host_name("kernel_cpy_f16_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<half, float>;
|
||||
template [[host_name("kernel_cpy_f16_f16")]] kernel kernel_cpy_t kernel_cpy_t_t<half, half>;
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_cpy_bf16_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<bfloat, float>;
|
||||
template [[host_name("kernel_cpy_bf16_bf16")]] kernel kernel_cpy_t kernel_cpy_t_t<bfloat, bfloat>;
|
||||
#endif
|
||||
|
||||
template<short QK,
|
||||
typename block_q,
|
||||
void (*quantize_func)(device const float *, device block_q &)>
|
||||
kernel void kernel_cpy_f32_q(
|
||||
constant ggml_metal_kargs_cpy & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig[2];
|
||||
const int32_t i02 = tgpig[1];
|
||||
const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
|
||||
const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
|
||||
|
||||
if (i01 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
|
||||
|
||||
const int32_t i3 = n / (args.ne2*args.ne1*args.ne0);
|
||||
const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0) / (args.ne1*args.ne0);
|
||||
const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0) / args.ne0;
|
||||
const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0)/QK;
|
||||
|
||||
device block_q * dst_data = (device block_q *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.nk0;) {
|
||||
device const float * src = (device const float *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + (i00*QK)*args.nb00);
|
||||
|
||||
quantize_func(src, dst_data[i00]);
|
||||
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_cpy_f32_q<QK8_0, block_q8_0, quantize_q8_0>) cpy_f_q_t;
|
||||
|
||||
template [[host_name("kernel_cpy_f32_q8_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK8_0, block_q8_0, quantize_q8_0>;
|
||||
template [[host_name("kernel_cpy_f32_q1_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK1_0, block_q1_0, quantize_q1_0>;
|
||||
template [[host_name("kernel_cpy_f32_q2_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK2_0, block_q2_0, quantize_q2_0>;
|
||||
template [[host_name("kernel_cpy_f32_q4_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK4_0, block_q4_0, quantize_q4_0>;
|
||||
template [[host_name("kernel_cpy_f32_q4_1")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK4_1, block_q4_1, quantize_q4_1>;
|
||||
template [[host_name("kernel_cpy_f32_q5_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK5_0, block_q5_0, quantize_q5_0>;
|
||||
template [[host_name("kernel_cpy_f32_q5_1")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK5_1, block_q5_1, quantize_q5_1>;
|
||||
template [[host_name("kernel_cpy_f32_iq4_nl")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK4_NL, block_iq4_nl, quantize_iq4_nl>;
|
||||
|
||||
template<typename T4x4, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread T4x4 &)>
|
||||
kernel void kernel_cpy_q_f32(
|
||||
constant ggml_metal_kargs_cpy & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig[2];
|
||||
const int32_t i02 = tgpig[1];
|
||||
const int32_t i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tpitg.y;
|
||||
const int32_t iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0;
|
||||
|
||||
if (i01 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00;
|
||||
|
||||
const int32_t i3 = n/(args.ne2*args.ne1*args.ne0);
|
||||
const int32_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0);
|
||||
const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
|
||||
const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
|
||||
|
||||
device const block_q * src_data = (device const block_q *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
|
||||
device T4x4 * dst_data = (device T4x4 *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.nk0;) {
|
||||
T4x4 temp;
|
||||
dequantize_func(src_data + i00/nl, i00%nl, temp);
|
||||
dst_data[i00] = temp;
|
||||
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>) cpy_q_f_t;
|
||||
|
||||
template [[host_name("kernel_cpy_q1_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q1_0, 8, dequantize_q1_0>;
|
||||
template [[host_name("kernel_cpy_q2_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q2_0, 4, dequantize_q2_0>;
|
||||
template [[host_name("kernel_cpy_q4_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>;
|
||||
template [[host_name("kernel_cpy_q4_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_1, 2, dequantize_q4_1>;
|
||||
template [[host_name("kernel_cpy_q5_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_0, 2, dequantize_q5_0>;
|
||||
template [[host_name("kernel_cpy_q5_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_1, 2, dequantize_q5_1>;
|
||||
template [[host_name("kernel_cpy_q8_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q8_0, 2, dequantize_q8_0>;
|
||||
|
||||
template [[host_name("kernel_cpy_q1_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q1_0, 8, dequantize_q1_0>;
|
||||
template [[host_name("kernel_cpy_q2_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q2_0, 4, dequantize_q2_0>;
|
||||
template [[host_name("kernel_cpy_q4_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_0, 2, dequantize_q4_0>;
|
||||
template [[host_name("kernel_cpy_q4_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_1, 2, dequantize_q4_1>;
|
||||
template [[host_name("kernel_cpy_q5_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_0, 2, dequantize_q5_0>;
|
||||
template [[host_name("kernel_cpy_q5_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_1, 2, dequantize_q5_1>;
|
||||
template [[host_name("kernel_cpy_q8_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q8_0, 2, dequantize_q8_0>;
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_concat(
|
||||
constant ggml_metal_kargs_concat & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = ntg.y == 1 ? tgpig.x : tgpig.x*ntg.y + tpitg.y;
|
||||
|
||||
if (i1 >= args.ne1) {
|
||||
return;
|
||||
}
|
||||
|
||||
int o[4] = {0, 0, 0, 0};
|
||||
o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03));
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
device const T * x;
|
||||
|
||||
if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
|
||||
x = (device const T *)(src0 + (i3 )*args.nb03 + (i2 )*args.nb02 + (i1 )*args.nb01 + (i0 )*args.nb00);
|
||||
} else {
|
||||
x = (device const T *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10);
|
||||
}
|
||||
|
||||
device T * y = (device T *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
*y = *x;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_concat<float>) kernel_concat_t;
|
||||
|
||||
template [[host_name("kernel_concat_f32")]] kernel kernel_concat_t kernel_concat<float>;
|
||||
template [[host_name("kernel_concat_f16")]] kernel kernel_concat_t kernel_concat<half>;
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_concat_bf16")]] kernel kernel_concat_t kernel_concat<bfloat>;
|
||||
#endif
|
||||
template [[host_name("kernel_concat_i8")]] kernel kernel_concat_t kernel_concat<char>;
|
||||
template [[host_name("kernel_concat_i16")]] kernel kernel_concat_t kernel_concat<short>;
|
||||
template [[host_name("kernel_concat_i32")]] kernel kernel_concat_t kernel_concat<int>;
|
||||
template [[host_name("kernel_concat_i64")]] kernel kernel_concat_t kernel_concat<long>;
|
||||
|
||||
template<typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread float4x4 &)>
|
||||
kernel void kernel_get_rows_q(
|
||||
constant ggml_metal_kargs_get_rows & args,
|
||||
device const void * src0,
|
||||
device const void * src1,
|
||||
device void * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const int32_t iw0 = tgpig.x/args.ne10;
|
||||
const int32_t i10 = tgpig.x%args.ne10;
|
||||
const int32_t i11 = tgpig.y;
|
||||
const int32_t i12 = tgpig.z;
|
||||
|
||||
const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0];
|
||||
|
||||
const int32_t i02 = i11;
|
||||
const int32_t i03 = i12;
|
||||
|
||||
auto psrc = (device const block_q *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 + r*args.nb01);
|
||||
auto pdst = (device float4x4 *) (( device char *) dst + i12*args.nb3 + i11*args.nb2 + i10*args.nb1);
|
||||
|
||||
for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) {
|
||||
float4x4 temp;
|
||||
dequantize_func(psrc + ind/nl, ind%nl, temp);
|
||||
pdst[ind] = temp;
|
||||
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
template<typename T0, typename T>
|
||||
kernel void kernel_get_rows_f(
|
||||
constant ggml_metal_kargs_get_rows & args,
|
||||
device const void * src0,
|
||||
device const void * src1,
|
||||
device void * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 ntg [[threads_per_threadgroup]]) {
|
||||
const int32_t iw0 = tgpig.x/args.ne10;
|
||||
const int32_t i10 = tgpig.x%args.ne10;
|
||||
const int32_t i11 = tgpig.y;
|
||||
const int32_t i12 = tgpig.z;
|
||||
|
||||
const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0];
|
||||
|
||||
const int32_t i02 = i11;
|
||||
const int32_t i03 = i12;
|
||||
|
||||
auto psrc = (const device T0 *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 + r*args.nb01);
|
||||
auto pdst = ( device T *) (( device char *) dst + i12*args.nb3 + i11*args.nb2 + i10*args.nb1);
|
||||
|
||||
for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) {
|
||||
pdst[ind] = psrc[ind];
|
||||
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
|
||||
|
||||
template [[host_name("kernel_get_rows_f32")]] kernel get_rows_f_t kernel_get_rows_f<float, float>;
|
||||
template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f<half, float>;
|
||||
template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f<bfloat, float>;
|
||||
#endif
|
||||
|
||||
typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
|
||||
|
||||
template [[host_name("kernel_get_rows_q1_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q1_0, 8, dequantize_q1_0>;
|
||||
template [[host_name("kernel_get_rows_q2_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q2_0, 4, dequantize_q2_0>;
|
||||
template [[host_name("kernel_get_rows_q4_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>;
|
||||
template [[host_name("kernel_get_rows_q4_1")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_1, 2, dequantize_q4_1>;
|
||||
template [[host_name("kernel_get_rows_q5_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q5_0, 2, dequantize_q5_0>;
|
||||
template [[host_name("kernel_get_rows_q5_1")]] kernel get_rows_q_t kernel_get_rows_q<block_q5_1, 2, dequantize_q5_1>;
|
||||
template [[host_name("kernel_get_rows_q8_0")]] kernel get_rows_q_t kernel_get_rows_q<block_q8_0, 2, dequantize_q8_0>;
|
||||
template [[host_name("kernel_get_rows_mxfp4")]] kernel get_rows_q_t kernel_get_rows_q<block_mxfp4, 2, dequantize_mxfp4>;
|
||||
template [[host_name("kernel_get_rows_q2_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q2_K, QK_NL, dequantize_q2_K>;
|
||||
template [[host_name("kernel_get_rows_q3_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q3_K, QK_NL, dequantize_q3_K>;
|
||||
template [[host_name("kernel_get_rows_q4_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q4_K, QK_NL, dequantize_q4_K>;
|
||||
template [[host_name("kernel_get_rows_q5_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q5_K, QK_NL, dequantize_q5_K>;
|
||||
template [[host_name("kernel_get_rows_q6_K")]] kernel get_rows_q_t kernel_get_rows_q<block_q6_K, QK_NL, dequantize_q6_K>;
|
||||
template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
|
||||
template [[host_name("kernel_get_rows_iq2_xs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xs, QK_NL, dequantize_iq2_xs>;
|
||||
template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_xxs, QK_NL, dequantize_iq3_xxs>;
|
||||
template [[host_name("kernel_get_rows_iq3_s")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_s, QK_NL, dequantize_iq3_s>;
|
||||
template [[host_name("kernel_get_rows_iq2_s")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_s, QK_NL, dequantize_iq2_s>;
|
||||
template [[host_name("kernel_get_rows_iq1_s")]] kernel get_rows_q_t kernel_get_rows_q<block_iq1_s, QK_NL, dequantize_iq1_s>;
|
||||
template [[host_name("kernel_get_rows_iq1_m")]] kernel get_rows_q_t kernel_get_rows_q<block_iq1_m, QK_NL, dequantize_iq1_m>;
|
||||
template [[host_name("kernel_get_rows_iq4_nl")]] kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl, 2, dequantize_iq4_nl>;
|
||||
template [[host_name("kernel_get_rows_iq4_xs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs, QK_NL, dequantize_iq4_xs>;
|
||||
|
||||
template<typename TS, typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
|
||||
kernel void kernel_set_rows_q32(
|
||||
constant ggml_metal_kargs_set_rows & args,
|
||||
device const void * src0,
|
||||
device const void * src1,
|
||||
device float * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint tiitg[[thread_index_in_threadgroup]],
|
||||
uint3 tptg [[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig.z;
|
||||
const int32_t i02 = tgpig.y;
|
||||
|
||||
const int32_t i12 = i03%args.ne12;
|
||||
const int32_t i11 = i02%args.ne11;
|
||||
|
||||
const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
|
||||
if (i01 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int32_t i10 = i01;
|
||||
const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
|
||||
|
||||
device block_q * dst_row = ( device block_q *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
||||
const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
||||
|
||||
for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
|
||||
quantize_func(src_row + 32*ind, dst_row[ind]);
|
||||
}
|
||||
}
|
||||
|
||||
template<typename TS, typename TI, typename TD>
|
||||
kernel void kernel_set_rows_f(
|
||||
constant ggml_metal_kargs_set_rows & args,
|
||||
device const void * src0,
|
||||
device const void * src1,
|
||||
device float * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint tiitg[[thread_index_in_threadgroup]],
|
||||
uint3 tptg [[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig.z;
|
||||
const int32_t i02 = tgpig.y;
|
||||
|
||||
const int32_t i12 = i03%args.ne12;
|
||||
const int32_t i11 = i02%args.ne11;
|
||||
|
||||
const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
|
||||
if (i01 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int32_t i10 = i01;
|
||||
const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
|
||||
|
||||
device TD * dst_row = ( device TD *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
||||
const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
||||
|
||||
for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
|
||||
dst_row[ind] = (TD) src_row[ind];
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_set_rows_f<float, int64_t, float>) set_rows_f_t;
|
||||
|
||||
template [[host_name("kernel_set_rows_f32_i64_f32")]] kernel set_rows_f_t kernel_set_rows_f<float, int64_t, float>;
|
||||
template [[host_name("kernel_set_rows_f32_i32_f32")]] kernel set_rows_f_t kernel_set_rows_f<float, int32_t, float>;
|
||||
template [[host_name("kernel_set_rows_f32_i64_f16")]] kernel set_rows_f_t kernel_set_rows_f<float, int64_t, half>;
|
||||
template [[host_name("kernel_set_rows_f32_i32_f16")]] kernel set_rows_f_t kernel_set_rows_f<float, int32_t, half>;
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_set_rows_f32_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f<float, int64_t, bfloat>;
|
||||
template [[host_name("kernel_set_rows_f32_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f<float, int32_t, bfloat>;
|
||||
#endif
|
||||
|
||||
template [[host_name("kernel_set_rows_f16_i64_f16")]] kernel set_rows_f_t kernel_set_rows_f<half, int64_t, half>;
|
||||
template [[host_name("kernel_set_rows_f16_i32_f16")]] kernel set_rows_f_t kernel_set_rows_f<half, int32_t, half>;
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_set_rows_bf16_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int64_t, bfloat>;
|
||||
template [[host_name("kernel_set_rows_bf16_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int32_t, bfloat>;
|
||||
#endif
|
||||
|
||||
typedef decltype(kernel_set_rows_q32<float, int64_t, block_q8_0, quantize_q8_0>) set_rows_q32_t;
|
||||
|
||||
template [[host_name("kernel_set_rows_f32_i64_q8_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q8_0, quantize_q8_0>;
|
||||
template [[host_name("kernel_set_rows_f32_i32_q8_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q8_0, quantize_q8_0>;
|
||||
template [[host_name("kernel_set_rows_f32_i64_q4_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_0, quantize_q4_0>;
|
||||
template [[host_name("kernel_set_rows_f32_i32_q4_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_0, quantize_q4_0>;
|
||||
template [[host_name("kernel_set_rows_f32_i64_q4_1")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_1, quantize_q4_1>;
|
||||
template [[host_name("kernel_set_rows_f32_i32_q4_1")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_1, quantize_q4_1>;
|
||||
template [[host_name("kernel_set_rows_f32_i64_q5_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_0, quantize_q5_0>;
|
||||
template [[host_name("kernel_set_rows_f32_i32_q5_0")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_0, quantize_q5_0>;
|
||||
template [[host_name("kernel_set_rows_f32_i64_q5_1")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_1, quantize_q5_1>;
|
||||
template [[host_name("kernel_set_rows_f32_i32_q5_1")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_1, quantize_q5_1>;
|
||||
template [[host_name("kernel_set_rows_f32_i64_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_iq4_nl, quantize_iq4_nl>;
|
||||
template [[host_name("kernel_set_rows_f32_i32_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_iq4_nl, quantize_iq4_nl>;
|
||||
|
||||
@@ -1,228 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
kernel void kernel_op_sum_f32(
|
||||
constant ggml_metal_kargs_sum & args,
|
||||
device const float * src0,
|
||||
device float * dst,
|
||||
threadgroup float * shmem_f32 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
if (args.np == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// TODO: become function constant
|
||||
const uint nsg = (ntg.x + 31) / 32;
|
||||
|
||||
float sumf = 0;
|
||||
|
||||
for (uint64_t i0 = tpitg.x; i0 < args.np; i0 += ntg.x) {
|
||||
sumf += src0[i0];
|
||||
}
|
||||
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_f32[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
float total = 0;
|
||||
|
||||
if (sgitg == 0) {
|
||||
float v = 0;
|
||||
|
||||
if (tpitg.x < nsg) {
|
||||
v = shmem_f32[tpitg.x];
|
||||
}
|
||||
|
||||
total = simd_sum(v);
|
||||
|
||||
if (tpitg.x == 0) {
|
||||
dst[0] = total;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
constant short FC_sum_rows_op [[function_constant(FC_SUM_ROWS + 0)]];
|
||||
|
||||
template <typename T0, typename T>
|
||||
kernel void kernel_sum_rows_impl(
|
||||
constant ggml_metal_kargs_sum_rows & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
#define FC_OP FC_sum_rows_op
|
||||
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = tgpig.x;
|
||||
|
||||
threadgroup T0 * shmem_t = (threadgroup T0 *) shmem;
|
||||
|
||||
if (sgitg == 0) {
|
||||
shmem_t[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
device const T0 * src_row = (device const T0 *) (src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03);
|
||||
device T * dst_row = (device T *) (dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3);
|
||||
|
||||
T0 sumf = T0(0.0f);
|
||||
|
||||
for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
|
||||
sumf += src_row[i0];
|
||||
}
|
||||
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shmem_t[sgitg] = sumf;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sumf = shmem_t[tiisg];
|
||||
sumf = simd_sum(sumf);
|
||||
|
||||
if (tpitg.x == 0) {
|
||||
if (FC_OP == OP_SUM_ROWS_NUM_MEAN) {
|
||||
if (is_same<float4, T0>::value) {
|
||||
dst_row[0] = sum(sumf) / (4*args.ne00);
|
||||
} else {
|
||||
dst_row[0] = sum(sumf) / args.ne00;
|
||||
}
|
||||
} else {
|
||||
dst_row[0] = sum(sumf);
|
||||
}
|
||||
}
|
||||
|
||||
#undef FC_OP
|
||||
}
|
||||
|
||||
typedef decltype(kernel_sum_rows_impl<float, float>) kernel_sum_rows_t;
|
||||
|
||||
template [[host_name("kernel_sum_rows_f32_f32")]] kernel kernel_sum_rows_t kernel_sum_rows_impl<float, float>;
|
||||
template [[host_name("kernel_sum_rows_f32_f32_4")]] kernel kernel_sum_rows_t kernel_sum_rows_impl<float4, float>;
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_cumsum_blk(
|
||||
constant ggml_metal_kargs_cumsum_blk & args,
|
||||
device const char * src0,
|
||||
device char * tmp,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int ib = tgpig[0]/args.ne01;
|
||||
|
||||
const int i00 = ib*ntg.x;
|
||||
const int i01 = tgpig[0]%args.ne01;
|
||||
const int i02 = tgpig[1];
|
||||
const int i03 = tgpig[2];
|
||||
|
||||
device const float * src0_row = (device const float *) (src0 +
|
||||
args.nb01*i01 +
|
||||
args.nb02*i02 +
|
||||
args.nb03*i03);
|
||||
|
||||
threadgroup float * shmem_f32 = (threadgroup float *) shmem;
|
||||
|
||||
float v = 0.0f;
|
||||
|
||||
if (i00 + tpitg.x < args.ne00) {
|
||||
v = src0_row[i00 + tpitg.x];
|
||||
}
|
||||
|
||||
float s = simd_prefix_inclusive_sum(v);
|
||||
|
||||
if (tiisg == N_SIMDWIDTH - 1) {
|
||||
shmem_f32[sgitg] = s;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (sgitg == 0) {
|
||||
shmem_f32[tiisg] = simd_prefix_exclusive_sum(shmem_f32[tiisg]);
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
s += shmem_f32[sgitg];
|
||||
|
||||
device float * dst_row = (device float *) dst +
|
||||
args.ne00*i01 +
|
||||
args.ne00*args.ne01*i02 +
|
||||
args.ne00*args.ne01*args.ne02*i03;
|
||||
|
||||
if (i00 + tpitg.x < args.ne00) {
|
||||
dst_row[i00 + tpitg.x] = s;
|
||||
}
|
||||
|
||||
if (args.outb && tpitg.x == ntg.x - 1) {
|
||||
device float * tmp_row = (device float *) tmp +
|
||||
args.net0*i01 +
|
||||
args.net0*args.net1*i02 +
|
||||
args.net0*args.net1*args.net2*i03;
|
||||
|
||||
tmp_row[ib] = s;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_cumsum_blk<float>) kernel_cumsum_blk_t;
|
||||
|
||||
template [[host_name("kernel_cumsum_blk_f32")]] kernel kernel_cumsum_blk_t kernel_cumsum_blk<float>;
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_cumsum_add(
|
||||
constant ggml_metal_kargs_cumsum_add & args,
|
||||
device const char * tmp,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int ib = tgpig[0]/args.ne01;
|
||||
|
||||
if (ib == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int i00 = ib*ntg.x;
|
||||
const int i01 = tgpig[0]%args.ne01;
|
||||
const int i02 = tgpig[1];
|
||||
const int i03 = tgpig[2];
|
||||
|
||||
device const float * tmp_row = (device const float *) (tmp +
|
||||
args.nbt1*i01 +
|
||||
args.nbt2*i02 +
|
||||
args.nbt3*i03);
|
||||
|
||||
device float * dst_row = (device float *) dst +
|
||||
args.ne00*i01 +
|
||||
args.ne00*args.ne01*i02 +
|
||||
args.ne00*args.ne01*args.ne02*i03;
|
||||
|
||||
if (i00 + tpitg.x < args.ne00) {
|
||||
dst_row[i00 + tpitg.x] += tmp_row[ib - 1];
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_cumsum_add<float>) kernel_cumsum_add_t;
|
||||
|
||||
template [[host_name("kernel_cumsum_add_f32")]] kernel kernel_cumsum_add_t kernel_cumsum_add<float>;
|
||||
@@ -1,318 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
constant bool FC_rope_is_imrope [[function_constant(FC_ROPE + 0)]];
|
||||
constant bool FC_rope_is_back [[function_constant(FC_ROPE + 1)]];
|
||||
|
||||
static float rope_yarn_ramp(const float low, const float high, const int i0) {
|
||||
const float y = (i0 / 2 - low) / max(0.001f, high - low);
|
||||
return 1.0f - min(1.0f, max(0.0f, y));
|
||||
}
|
||||
|
||||
// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn
|
||||
// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng.
|
||||
static void rope_yarn(
|
||||
float theta_extrap, float freq_scale, float corr_dims[2], int i0, float ext_factor, float mscale,
|
||||
thread float * cos_theta, thread float * sin_theta) {
|
||||
// Get n-d rotational scaling corrected for extrapolation
|
||||
float theta_interp = freq_scale * theta_extrap;
|
||||
float theta = theta_interp;
|
||||
if (ext_factor != 0.0f) {
|
||||
float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * ext_factor;
|
||||
theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix;
|
||||
|
||||
// Get n-d magnitude scaling corrected for interpolation
|
||||
mscale *= 1.0f + 0.1f * log(1.0f / freq_scale);
|
||||
}
|
||||
*cos_theta = cos(theta) * mscale;
|
||||
*sin_theta = sin(theta) * mscale;
|
||||
if (FC_rope_is_back) {
|
||||
*sin_theta *= -1.0f;
|
||||
}
|
||||
}
|
||||
|
||||
// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get
|
||||
// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))`
|
||||
static float rope_yarn_corr_factor(int n_dims, int n_ctx_orig, float n_rot, float base) {
|
||||
return n_dims * log(n_ctx_orig / (n_rot * 2 * M_PI_F)) / (2 * log(base));
|
||||
}
|
||||
|
||||
static void rope_yarn_corr_dims(
|
||||
int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]
|
||||
) {
|
||||
// start and end correction dims
|
||||
dims[0] = max(0.0f, floor(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_fast, freq_base)));
|
||||
dims[1] = min(n_dims - 1.0f, ceil(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_slow, freq_base)));
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_rope_norm(
|
||||
constant ggml_metal_kargs_rope & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 tptg [[threads_per_threadgroup]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]]) {
|
||||
const int i3 = tgpig[2];
|
||||
const int i2 = tgpig[1];
|
||||
const int i1 = tgpig[0];
|
||||
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
||||
|
||||
device const int32_t * pos = (device const int32_t *) src1;
|
||||
|
||||
const float theta_base = (float) pos[i2];
|
||||
const float inv_ndims = -1.f/args.n_dims;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
|
||||
for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
|
||||
if (i0 < args.n_dims) {
|
||||
const int ic = i0/2;
|
||||
|
||||
const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
|
||||
|
||||
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
||||
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
|
||||
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
|
||||
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
const float x0 = src[0];
|
||||
const float x1 = src[1];
|
||||
|
||||
dst_data[0] = x0*cos_theta - x1*sin_theta;
|
||||
dst_data[1] = x0*sin_theta + x1*cos_theta;
|
||||
} else {
|
||||
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
|
||||
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
dst_data[0] = src[0];
|
||||
dst_data[1] = src[1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_rope_neox(
|
||||
constant ggml_metal_kargs_rope & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 tptg [[threads_per_threadgroup]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]]) {
|
||||
const int i3 = tgpig[2];
|
||||
const int i2 = tgpig[1];
|
||||
const int i1 = tgpig[0];
|
||||
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
||||
|
||||
device const int32_t * pos = (device const int32_t *) src1;
|
||||
|
||||
const float theta_base = (float) pos[i2];
|
||||
const float inv_ndims = -1.f/args.n_dims;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
|
||||
for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
|
||||
if (i0 < args.n_dims) {
|
||||
const int ic = i0/2;
|
||||
|
||||
const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
|
||||
|
||||
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
||||
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
|
||||
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
|
||||
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + ic*args.nb0);
|
||||
|
||||
const float x0 = src[0];
|
||||
const float x1 = src[args.n_dims/2];
|
||||
|
||||
dst_data[0] = x0*cos_theta - x1*sin_theta;
|
||||
dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
|
||||
} else {
|
||||
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
|
||||
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
dst_data[0] = src[0];
|
||||
dst_data[1] = src[1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_rope_multi(
|
||||
constant ggml_metal_kargs_rope & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 tptg [[threads_per_threadgroup]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]]) {
|
||||
const int i3 = tgpig[2];
|
||||
const int i2 = tgpig[1];
|
||||
const int i1 = tgpig[0];
|
||||
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
||||
|
||||
device const int32_t * pos = (device const int32_t *) src1;
|
||||
|
||||
const float inv_ndims = -1.f/args.n_dims;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
|
||||
for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
|
||||
if (i0 < args.n_dims) {
|
||||
const int ic = i0/2;
|
||||
|
||||
// mrope theta calculations
|
||||
// note: the rest is the same as kernel_rope_neox
|
||||
const int sect_dims = args.sect_0 + args.sect_1 + args.sect_2 + args.sect_3;
|
||||
const int sec_w01 = args.sect_0 + args.sect_1; // end of section 1
|
||||
const int sec_w012 = args.sect_0 + args.sect_1 + args.sect_2; // end of section 2
|
||||
const int sector = ic % sect_dims;
|
||||
|
||||
float theta_base;
|
||||
if (FC_rope_is_imrope) {
|
||||
if (sector % 3 == 1 && sector < 3 * args.sect_1) { // h
|
||||
theta_base = (float) pos[i2 + args.ne02 * 1];
|
||||
} else if (sector % 3 == 2 && sector < 3 * args.sect_2) { // w
|
||||
theta_base = (float) pos[i2 + args.ne02 * 2];
|
||||
} else if (sector % 3 == 0 && sector < 3 * args.sect_0) { // t
|
||||
theta_base = (float) pos[i2 + args.ne02 * 0];
|
||||
} else { // e
|
||||
theta_base = (float) pos[i2 + args.ne02 * 3];
|
||||
}
|
||||
} else {
|
||||
if (sector < args.sect_0) {
|
||||
theta_base = (float) pos[i2];
|
||||
} else if (sector < sec_w01) {
|
||||
theta_base = (float) pos[i2 + args.ne02 * 1];
|
||||
} else if (sector < sec_w012) {
|
||||
theta_base = (float) pos[i2 + args.ne02 * 2];
|
||||
} else {
|
||||
theta_base = (float) pos[i2 + args.ne02 * 3];
|
||||
}
|
||||
}
|
||||
// end of mrope
|
||||
|
||||
const float theta = theta_base * pow(args.freq_base, inv_ndims*i0);
|
||||
|
||||
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
||||
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
|
||||
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
|
||||
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + ic*args.nb0);
|
||||
|
||||
const float x0 = src[0];
|
||||
const float x1 = src[args.n_dims/2];
|
||||
|
||||
dst_data[0] = x0*cos_theta - x1*sin_theta;
|
||||
dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta;
|
||||
} else {
|
||||
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
|
||||
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
dst_data[0] = src[0];
|
||||
dst_data[1] = src[1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_rope_vision(
|
||||
constant ggml_metal_kargs_rope & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 tptg [[threads_per_threadgroup]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]]) {
|
||||
const int i3 = tgpig[2];
|
||||
const int i2 = tgpig[1];
|
||||
const int i1 = tgpig[0];
|
||||
|
||||
float corr_dims[2];
|
||||
rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims);
|
||||
|
||||
device const int32_t * pos = (device const int32_t *) src1;
|
||||
|
||||
const float inv_ndims = -1.f/args.n_dims;
|
||||
|
||||
float cos_theta;
|
||||
float sin_theta;
|
||||
|
||||
for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) {
|
||||
if (i0 < 2*args.n_dims) { // different from kernel_rope_multi
|
||||
const int ic = i0/2;
|
||||
|
||||
// mrope theta calculations (only support 2 dimensions)
|
||||
const int sect_dims = args.sect_0 + args.sect_1;
|
||||
const int sector = ic % sect_dims;
|
||||
|
||||
float p;
|
||||
float theta_base;
|
||||
if (sector < args.sect_1) {
|
||||
p = (float) sector;
|
||||
theta_base = (float) pos[i2];
|
||||
} else {
|
||||
p = (float) sector - args.sect_0;
|
||||
theta_base = (float) pos[i2 + args.ne02];
|
||||
}
|
||||
|
||||
const float theta = theta_base * pow(args.freq_base, 2.0f * inv_ndims * p);
|
||||
// end of mrope
|
||||
|
||||
const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f;
|
||||
|
||||
rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta);
|
||||
|
||||
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00);
|
||||
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + ic*args.nb0);
|
||||
|
||||
const float x0 = src[0];
|
||||
const float x1 = src[args.n_dims]; // different from kernel_rope_multi
|
||||
|
||||
dst_data[0] = x0*cos_theta - x1*sin_theta;
|
||||
dst_data[args.n_dims] = x0*sin_theta + x1*cos_theta; // different from kernel_rope_multi
|
||||
} else {
|
||||
device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00);
|
||||
device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
dst_data[0] = src[0];
|
||||
dst_data[1] = src[1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_rope_norm<float>) kernel_rope_norm_t;
|
||||
typedef decltype(kernel_rope_neox<float>) kernel_rope_neox_t;
|
||||
typedef decltype(kernel_rope_multi<float>) kernel_rope_multi_t;
|
||||
typedef decltype(kernel_rope_vision<float>) kernel_rope_vision_t;
|
||||
|
||||
template [[host_name("kernel_rope_norm_f32")]] kernel kernel_rope_norm_t kernel_rope_norm<float>;
|
||||
template [[host_name("kernel_rope_norm_f16")]] kernel kernel_rope_norm_t kernel_rope_norm<half>;
|
||||
|
||||
template [[host_name("kernel_rope_neox_f32")]] kernel kernel_rope_neox_t kernel_rope_neox<float>;
|
||||
template [[host_name("kernel_rope_neox_f16")]] kernel kernel_rope_neox_t kernel_rope_neox<half>;
|
||||
|
||||
template [[host_name("kernel_rope_multi_f32")]] kernel kernel_rope_multi_t kernel_rope_multi<float>;
|
||||
template [[host_name("kernel_rope_multi_f16")]] kernel kernel_rope_multi_t kernel_rope_multi<half>;
|
||||
|
||||
template [[host_name("kernel_rope_vision_f32")]] kernel kernel_rope_vision_t kernel_rope_vision<float>;
|
||||
template [[host_name("kernel_rope_vision_f16")]] kernel kernel_rope_vision_t kernel_rope_vision<half>;
|
||||
@@ -1,223 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_soft_max(
|
||||
constant ggml_metal_kargs_soft_max & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
threadgroup float * buf [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint sgitg[[simdgroup_index_in_threadgroup]],
|
||||
uint tiisg[[thread_index_in_simdgroup]],
|
||||
uint3 tptg[[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig.z;
|
||||
const int32_t i02 = tgpig.y;
|
||||
const int32_t i01 = tgpig.x;
|
||||
|
||||
const int32_t i13 = i03%args.ne13;
|
||||
const int32_t i12 = i02%args.ne12;
|
||||
const int32_t i11 = i01;
|
||||
|
||||
device const float * psrc0 = (device const float *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
||||
device const T * pmask = src1 != src0 ? (device const T * ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
|
||||
device const float * psrc2 = src2 != src0 ? (device const float *) (src2) : nullptr;
|
||||
device float * pdst = (device float *) (dst + i01*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
||||
|
||||
float slope = 1.0f;
|
||||
|
||||
// ALiBi
|
||||
if (args.max_bias > 0.0f) {
|
||||
const int32_t h = i02;
|
||||
|
||||
const float base = h < args.n_head_log2 ? args.m0 : args.m1;
|
||||
const int exp = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
|
||||
|
||||
slope = pow(base, exp);
|
||||
}
|
||||
|
||||
// parallel max
|
||||
float lmax = psrc2 ? psrc2[i02] : -INFINITY;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
||||
lmax = MAX(lmax, psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f));
|
||||
}
|
||||
|
||||
// find the max value in the block
|
||||
float max_val = simd_max(lmax);
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = -INFINITY;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = max_val;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
max_val = buf[tiisg];
|
||||
max_val = simd_max(max_val);
|
||||
}
|
||||
|
||||
// parallel sum
|
||||
float lsum = 0.0f;
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
||||
const float exp_psrc0 = exp((psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f)) - max_val);
|
||||
lsum += exp_psrc0;
|
||||
pdst[i00] = exp_psrc0;
|
||||
}
|
||||
|
||||
// This barrier fixes a failing test
|
||||
// ref: https://github.com/ggml-org/ggml/pull/621#discussion_r1425156335
|
||||
threadgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
float sum = simd_sum(lsum);
|
||||
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = sum;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sum = buf[tiisg];
|
||||
sum = simd_sum(sum);
|
||||
}
|
||||
|
||||
if (psrc2) {
|
||||
sum += exp(psrc2[i02] - max_val);
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f/sum;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
|
||||
pdst[i00] *= inv_sum;
|
||||
}
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_soft_max_4(
|
||||
constant ggml_metal_kargs_soft_max & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device const char * src2,
|
||||
device char * dst,
|
||||
threadgroup float * buf [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint sgitg[[simdgroup_index_in_threadgroup]],
|
||||
uint tiisg[[thread_index_in_simdgroup]],
|
||||
uint3 tptg[[threads_per_threadgroup]]) {
|
||||
const int32_t i03 = tgpig.z;
|
||||
const int32_t i02 = tgpig.y;
|
||||
const int32_t i01 = tgpig.x;
|
||||
|
||||
const int32_t i13 = i03%args.ne13;
|
||||
const int32_t i12 = i02%args.ne12;
|
||||
const int32_t i11 = i01;
|
||||
|
||||
device const float4 * psrc4 = (device const float4 *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
|
||||
device const T * pmask = src1 != src0 ? (device const T * ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
|
||||
device const float * psrc2 = src2 != src0 ? (device const float * ) (src2) : nullptr;
|
||||
device float4 * pdst4 = (device float4 *) (dst + i01*args.nb1 + i02*args.nb2 + i03*args.nb3);
|
||||
|
||||
float slope = 1.0f;
|
||||
|
||||
if (args.max_bias > 0.0f) {
|
||||
const int32_t h = i02;
|
||||
|
||||
const float base = h < args.n_head_log2 ? args.m0 : args.m1;
|
||||
const int exp = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
|
||||
|
||||
slope = pow(base, exp);
|
||||
}
|
||||
|
||||
// parallel max
|
||||
float4 lmax4 = psrc2 ? psrc2[i02] : -INFINITY;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
||||
lmax4 = fmax(lmax4, psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f)));
|
||||
}
|
||||
|
||||
const float lmax = MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
|
||||
|
||||
float max_val = simd_max(lmax);
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = -INFINITY;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = max_val;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
max_val = buf[tiisg];
|
||||
max_val = simd_max(max_val);
|
||||
}
|
||||
|
||||
// parallel sum
|
||||
float4 lsum4 = 0.0f;
|
||||
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
||||
const float4 exp_psrc4 = exp((psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f))) - max_val);
|
||||
lsum4 += exp_psrc4;
|
||||
pdst4[i00] = exp_psrc4;
|
||||
}
|
||||
|
||||
const float lsum = lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
|
||||
|
||||
// This barrier fixes a failing test
|
||||
// ref: https://github.com/ggml-org/ggml/pull/621#discussion_r1425156335
|
||||
threadgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
float sum = simd_sum(lsum);
|
||||
|
||||
if (tptg.x > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
buf[tiisg] = 0.0f;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
buf[sgitg] = sum;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sum = buf[tiisg];
|
||||
sum = simd_sum(sum);
|
||||
}
|
||||
|
||||
if (psrc2) {
|
||||
sum += exp(psrc2[i02] - max_val);
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f/sum;
|
||||
|
||||
for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
|
||||
pdst4[i00] *= inv_sum;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_soft_max<float>) kernel_soft_max_t;
|
||||
typedef decltype(kernel_soft_max_4<float4>) kernel_soft_max_4_t;
|
||||
|
||||
template [[host_name("kernel_soft_max_f16")]] kernel kernel_soft_max_t kernel_soft_max<half>;
|
||||
template [[host_name("kernel_soft_max_f32")]] kernel kernel_soft_max_t kernel_soft_max<float>;
|
||||
template [[host_name("kernel_soft_max_f16_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<half4>;
|
||||
template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<float4>;
|
||||
@@ -1,75 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
constant short FC_solve_tri_nsg [[function_constant(FC_SOLVE_TRI + 0)]];
|
||||
constant short FC_solve_tri_n [[function_constant(FC_SOLVE_TRI + 1)]];
|
||||
constant short FC_solve_tri_k [[function_constant(FC_SOLVE_TRI + 2)]];
|
||||
|
||||
kernel void kernel_solve_tri_f32(
|
||||
constant ggml_metal_kargs_solve_tri & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
threadgroup char * shmem [[threadgroup(0)]],
|
||||
ushort3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
|
||||
const short NSG = FC_solve_tri_nsg;
|
||||
const short N = FC_solve_tri_n;
|
||||
const short K = FC_solve_tri_k;
|
||||
const short NP = PAD2(N, NW);
|
||||
|
||||
const int32_t i03 = tgpig.z;
|
||||
const int32_t i02 = tgpig.y;
|
||||
const int32_t i01 = tgpig.x*NSG + sgitg;
|
||||
|
||||
threadgroup float * sh0 = (threadgroup float *) shmem;
|
||||
|
||||
device const float * src0_ptr = (device const float *)(src0 + i02 * args.nb02 + i03 * args.nb03) + sgitg*N;
|
||||
device const float * src1_ptr = (device const float *)(src1 + i02 * args.nb12 + i03 * args.nb13) + i01;
|
||||
device float * dst_ptr = (device float *)(dst + i02 * args.nb2 + i03 * args.nb3) + i01;
|
||||
|
||||
for (short rr = 0; rr < N; rr += NSG) {
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
{
|
||||
threadgroup float * sh0_cur = sh0 + sgitg*NP;
|
||||
|
||||
for (short t = 0; t*NW < N; ++t) {
|
||||
const short idx = t*NW + tiisg;
|
||||
sh0_cur[idx] = src0_ptr[idx];
|
||||
}
|
||||
|
||||
src0_ptr += NSG*N;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (i01 >= args.ne10) {
|
||||
continue;
|
||||
}
|
||||
|
||||
for (short ir = 0; ir < NSG && rr + ir < N; ++ir) {
|
||||
const short r = rr + ir;
|
||||
|
||||
threadgroup float * sh0_cur = sh0 + ir*NP;
|
||||
|
||||
float sum = 0.0f;
|
||||
|
||||
for (short t = 0; t*NW < r; ++t) {
|
||||
const short idx = t*NW + tiisg;
|
||||
sum += sh0_cur[idx] * dst_ptr[idx*K] * (idx < r);
|
||||
}
|
||||
|
||||
sum = simd_sum(sum);
|
||||
|
||||
if (tiisg == 0) {
|
||||
const float diag = sh0_cur[r];
|
||||
|
||||
dst_ptr[r*K] = (src1_ptr[r*K] - sum) / diag;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,279 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
// ref: ggml.c:ggml_compute_forward_ssm_conv_f32
|
||||
kernel void kernel_ssm_conv_f32_f32(
|
||||
constant ggml_metal_kargs_ssm_conv & args,
|
||||
device const void * src0,
|
||||
device const void * src1,
|
||||
device float * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
const int64_t ir = tgpig.x;
|
||||
const int64_t i2 = tgpig.y;
|
||||
const int64_t i3 = tgpig.z;
|
||||
|
||||
const int64_t nc = args.ne10;
|
||||
//const int64_t ncs = args.ne00;
|
||||
//const int64_t nr = args.ne01;
|
||||
//const int64_t n_t = args.ne1;
|
||||
//const int64_t n_s = args.ne2;
|
||||
|
||||
device const float * s = (device const float *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
|
||||
device const float * c = (device const float *) ((device const char *) src1 + ir*args.nb11);
|
||||
device float * x = (device float *) ((device char *) dst + ir*args.nb0 + i2*args.nb1 + i3*args.nb2);
|
||||
|
||||
float sumf = 0.0f;
|
||||
|
||||
for (int64_t i0 = 0; i0 < nc; ++i0) {
|
||||
sumf += s[i0] * c[i0];
|
||||
}
|
||||
|
||||
x[0] = sumf;
|
||||
}
|
||||
|
||||
kernel void kernel_ssm_conv_f32_f32_4(
|
||||
constant ggml_metal_kargs_ssm_conv & args,
|
||||
device const void * src0,
|
||||
device const void * src1,
|
||||
device float * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
const int64_t ir = tgpig.x;
|
||||
const int64_t i2 = tgpig.y;
|
||||
const int64_t i3 = tgpig.z;
|
||||
|
||||
const int64_t nc = args.ne10;
|
||||
//const int64_t ncs = args.ne00;
|
||||
//const int64_t nr = args.ne01;
|
||||
//const int64_t n_t = args.ne1;
|
||||
//const int64_t n_s = args.ne2;
|
||||
|
||||
device const float4 * s = (device const float4 *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
|
||||
device const float4 * c = (device const float4 *) ((device const char *) src1 + ir*args.nb11);
|
||||
device float * x = (device float *) ((device char *) dst + ir*args.nb0 + i2*args.nb1 + i3*args.nb2);
|
||||
|
||||
float sumf = 0.0f;
|
||||
|
||||
for (int64_t i0 = 0; i0 < nc/4; ++i0) {
|
||||
sumf += dot(s[i0], c[i0]);
|
||||
}
|
||||
|
||||
x[0] = sumf;
|
||||
}
|
||||
|
||||
constant short FC_ssm_conv_bs [[function_constant(FC_SSM_CONV + 0)]];
|
||||
|
||||
// Batched version: each threadgroup processes multiple tokens for better efficiency
|
||||
// Thread layout: each thread handles one token, threadgroup covers BATCH_SIZE tokens
|
||||
kernel void kernel_ssm_conv_f32_f32_batched(
|
||||
constant ggml_metal_kargs_ssm_conv & args,
|
||||
device const void * src0,
|
||||
device const void * src1,
|
||||
device float * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
// tgpig.x = row index (ir)
|
||||
// tgpig.y = batch of tokens (i2_base / BATCH_SIZE)
|
||||
// tgpig.z = sequence index (i3)
|
||||
// tpitg.x = thread within batch (0..BATCH_SIZE-1)
|
||||
const short BATCH_SIZE = FC_ssm_conv_bs;
|
||||
|
||||
const int64_t ir = tgpig.x;
|
||||
const int64_t i2_base = tgpig.y * BATCH_SIZE;
|
||||
const int64_t i3 = tgpig.z;
|
||||
const int64_t i2_off = tpitg.x;
|
||||
const int64_t i2 = i2_base + i2_off;
|
||||
|
||||
const int64_t nc = args.ne10; // conv kernel size (typically 4)
|
||||
const int64_t n_t = args.ne1; // number of tokens
|
||||
|
||||
// Bounds check for partial batches at the end
|
||||
if (i2 >= n_t) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Load conv weights (shared across all tokens for this row)
|
||||
device const float * c = (device const float *) ((device const char *) src1 + ir*args.nb11);
|
||||
|
||||
// Load source for this specific token
|
||||
device const float * s = (device const float *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
|
||||
|
||||
// Output location for this token
|
||||
device float * x = (device float *) ((device char *) dst + ir*args.nb0 + i2*args.nb1 + i3*args.nb2);
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int64_t i0 = 0; i0 < nc; ++i0) {
|
||||
sumf += s[i0] * c[i0];
|
||||
}
|
||||
|
||||
x[0] = sumf;
|
||||
}
|
||||
|
||||
kernel void kernel_ssm_conv_f32_f32_batched_4(
|
||||
constant ggml_metal_kargs_ssm_conv & args,
|
||||
device const void * src0,
|
||||
device const void * src1,
|
||||
device float * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
// tgpig.x = row index (ir)
|
||||
// tgpig.y = batch of tokens (i2_base / BATCH_SIZE)
|
||||
// tgpig.z = sequence index (i3)
|
||||
// tpitg.x = thread within batch (0..BATCH_SIZE-1)
|
||||
const short BATCH_SIZE = FC_ssm_conv_bs;
|
||||
|
||||
const int64_t ir = tgpig.x;
|
||||
const int64_t i2_base = tgpig.y * BATCH_SIZE;
|
||||
const int64_t i3 = tgpig.z;
|
||||
const int64_t i2_off = tpitg.x;
|
||||
const int64_t i2 = i2_base + i2_off;
|
||||
|
||||
const int64_t nc = args.ne10; // conv kernel size (typically 4)
|
||||
const int64_t n_t = args.ne1; // number of tokens
|
||||
|
||||
// Bounds check for partial batches at the end
|
||||
if (i2 >= n_t) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Load conv weights (shared across all tokens for this row)
|
||||
device const float4 * c = (device const float4 *) ((device const char *) src1 + ir*args.nb11);
|
||||
|
||||
// Load source for this specific token
|
||||
device const float4 * s = (device const float4 *) ((device const char *) src0 + ir*args.nb01 + i2*args.nb00 + i3*args.nb02);
|
||||
|
||||
// Output location for this token
|
||||
device float * x = (device float *) ((device char *) dst + ir*args.nb0 + i2*args.nb1 + i3*args.nb2);
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int64_t i0 = 0; i0 < nc/4; ++i0) {
|
||||
sumf += dot(s[i0], c[i0]);
|
||||
}
|
||||
|
||||
x[0] = sumf;
|
||||
}
|
||||
|
||||
// ref: ggml.c:ggml_compute_forward_ssm_scan_f32, Mamba-2 part
|
||||
// Optimized version: reduces redundant memory loads by having one thread load shared values
|
||||
kernel void kernel_ssm_scan_f32(
|
||||
constant ggml_metal_kargs_ssm_scan & args,
|
||||
device const void * src0,
|
||||
device const void * src1,
|
||||
device const void * src2,
|
||||
device const void * src3,
|
||||
device const void * src4,
|
||||
device const void * src5,
|
||||
device const void * src6,
|
||||
device float * dst,
|
||||
threadgroup float * shared [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgptg[[simdgroups_per_threadgroup]],
|
||||
uint3 tgpg[[threadgroups_per_grid]]) {
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
|
||||
// Shared memory layout:
|
||||
// [0..sgptg*NW-1]: partial sums for reduction (existing)
|
||||
// [sgptg*NW..sgptg*NW+sgptg-1]: pre-computed x_dt values for each token in batch
|
||||
// [sgptg*NW+sgptg..sgptg*NW+2*sgptg-1]: pre-computed dA values for each token in batch
|
||||
threadgroup float * shared_sums = shared;
|
||||
threadgroup float * shared_x_dt = shared + sgptg * NW;
|
||||
threadgroup float * shared_dA = shared + sgptg * NW + sgptg;
|
||||
|
||||
shared_sums[tpitg.x] = 0.0f;
|
||||
|
||||
const int32_t i0 = tpitg.x;
|
||||
const int32_t i1 = tgpig.x;
|
||||
const int32_t ir = tgpig.y; // current head
|
||||
const int32_t i3 = tgpig.z; // current seq
|
||||
|
||||
const int32_t nc = args.d_state;
|
||||
const int32_t nr = args.d_inner;
|
||||
const int32_t nh = args.n_head;
|
||||
const int32_t ng = args.n_group;
|
||||
const int32_t n_t = args.n_seq_tokens;
|
||||
|
||||
const int32_t s_off = args.s_off;
|
||||
|
||||
device const int32_t * ids = (device const int32_t *) src6;
|
||||
|
||||
device const float * s0_buff = (device const float *) ((device const char *) src0 + ir*args.nb02 + ids[i3]*args.nb03);
|
||||
device float * s_buff = (device float *) ((device char *) dst + ir*args.nb02 + i3*args.nb03 + s_off);
|
||||
|
||||
const int32_t i = i0 + i1*nc;
|
||||
const int32_t g = ir / (nh / ng); // repeat_interleave
|
||||
|
||||
float s0 = s0_buff[i];
|
||||
float s = 0.0f;
|
||||
|
||||
device const float * A = (device const float *) ((device const char *) src3 + ir*args.nb31); // {ne30, nh}
|
||||
|
||||
const float A0 = A[i0%args.ne30];
|
||||
|
||||
device const float * x = (device const float *)((device const char *) src1 + i1*args.nb10 + ir*args.nb11 + i3*args.nb13); // {dim, nh, nt, ns}
|
||||
device const float * dt = (device const float *)((device const char *) src2 + ir*args.nb20 + i3*args.nb22); // {nh, nt, ns}
|
||||
device const float * B = (device const float *)((device const char *) src4 + g*args.nb41 + i3*args.nb43); // {d_state, ng, nt, ns}
|
||||
device const float * C = (device const float *)((device const char *) src5 + g*args.nb51 + i3*args.nb53); // {d_state, ng, nt, ns}
|
||||
|
||||
device float * y = dst + (i1 + ir*(nr) + i3*(n_t*nh*nr)); // {dim, nh, nt, ns}
|
||||
|
||||
for (int i2 = 0; i2 < n_t; i2 += sgptg) {
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Pre-compute x_dt and dA for this batch of tokens
|
||||
// Only first sgptg threads do the loads and expensive math
|
||||
if (i0 < sgptg && i2 + i0 < n_t) {
|
||||
// ns12 and ns21 are element strides (nb12/nb10, nb21/nb20)
|
||||
device const float * x_t = x + i0 * args.ns12;
|
||||
device const float * dt_t = dt + i0 * args.ns21;
|
||||
|
||||
const float dt0 = dt_t[0];
|
||||
const float dtsp = dt0 <= 20.0f ? log(1.0f + exp(dt0)) : dt0;
|
||||
shared_x_dt[i0] = x_t[0] * dtsp;
|
||||
shared_dA[i0] = dtsp; // Store dtsp, compute exp(dtsp * A0) per-thread since A0 varies
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (int t = 0; t < sgptg && i2 + t < n_t; t++) {
|
||||
const float x_dt = shared_x_dt[t];
|
||||
const float dA = exp(shared_dA[t] * A0);
|
||||
|
||||
s = (s0 * dA) + (B[i0] * x_dt);
|
||||
|
||||
const float sumf = simd_sum(s * C[i0]);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shared_sums[t*NW + sgitg] = sumf;
|
||||
}
|
||||
|
||||
// recurse
|
||||
s0 = s;
|
||||
|
||||
B += args.ns42;
|
||||
C += args.ns52;
|
||||
}
|
||||
|
||||
// Advance pointers for next batch
|
||||
x += sgptg * args.ns12;
|
||||
dt += sgptg * args.ns21;
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
const float sumf = simd_sum(shared_sums[sgitg*NW + tiisg]);
|
||||
|
||||
if (tiisg == 0 && i2 + sgitg < n_t) {
|
||||
y[sgitg*nh*nr] = sumf;
|
||||
}
|
||||
|
||||
y += sgptg*nh*nr;
|
||||
}
|
||||
|
||||
s_buff[i] = s;
|
||||
}
|
||||
@@ -1,69 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
template<uint32_t ttype>
|
||||
bool _ggml_vec_tri_cmp(const int i, const int r);
|
||||
|
||||
template<>
|
||||
bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_LOWER */ 3>(const int i, const int r) {
|
||||
return i < r;
|
||||
}
|
||||
|
||||
template<>
|
||||
bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_LOWER_DIAG */ 2>(const int i, const int r) {
|
||||
return i <= r;
|
||||
}
|
||||
|
||||
template<>
|
||||
bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_UPPER */ 1>(const int i, const int r) {
|
||||
return i > r;
|
||||
}
|
||||
|
||||
template<>
|
||||
bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_UPPER_DIAG */ 0>(const int i, const int r) {
|
||||
return i >= r;
|
||||
}
|
||||
|
||||
template<typename T, int ttype>
|
||||
kernel void kernel_tri(
|
||||
constant ggml_metal_kargs_tri & args,
|
||||
device const char * src0,
|
||||
device const char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
const int i3 = tgpig.z;
|
||||
const int i2 = tgpig.y;
|
||||
const int i1 = tgpig.x;
|
||||
|
||||
if (i3 >= args.ne03 || i2 >= args.ne02 || i1 >= args.ne01) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const T * src_row = (device const T *) ((device const char *) src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03);
|
||||
device T * dst_row = (device T *) ((device char *) dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3);
|
||||
|
||||
// Each thread is a single element of the row if ne00 < max threads per
|
||||
// threadgroup, so this will loop once for each index that this thread is
|
||||
// responsible for
|
||||
for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
|
||||
// Use the comparison as a mask for branchless
|
||||
dst_row[i0] = static_cast<T>(_ggml_vec_tri_cmp<ttype>(i0, i1)) * src_row[i0];
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_tri<float, 0>) kernel_tri_t;
|
||||
|
||||
template [[host_name("kernel_tri_f32_0")]] kernel kernel_tri_t kernel_tri<float, 0>;
|
||||
template [[host_name("kernel_tri_f32_1")]] kernel kernel_tri_t kernel_tri<float, 1>;
|
||||
template [[host_name("kernel_tri_f32_2")]] kernel kernel_tri_t kernel_tri<float, 2>;
|
||||
template [[host_name("kernel_tri_f32_3")]] kernel kernel_tri_t kernel_tri<float, 3>;
|
||||
template [[host_name("kernel_tri_f16_0")]] kernel kernel_tri_t kernel_tri<half, 0>;
|
||||
template [[host_name("kernel_tri_f16_1")]] kernel kernel_tri_t kernel_tri<half, 1>;
|
||||
template [[host_name("kernel_tri_f16_2")]] kernel kernel_tri_t kernel_tri<half, 2>;
|
||||
template [[host_name("kernel_tri_f16_3")]] kernel kernel_tri_t kernel_tri<half, 3>;
|
||||
#if defined(GGML_METAL_HAS_BF16)
|
||||
template [[host_name("kernel_tri_bf16_0")]] kernel kernel_tri_t kernel_tri<bfloat, 0>;
|
||||
template [[host_name("kernel_tri_bf16_1")]] kernel kernel_tri_t kernel_tri<bfloat, 1>;
|
||||
template [[host_name("kernel_tri_bf16_2")]] kernel kernel_tri_t kernel_tri<bfloat, 2>;
|
||||
template [[host_name("kernel_tri_bf16_3")]] kernel kernel_tri_t kernel_tri<bfloat, 3>;
|
||||
#endif
|
||||
@@ -1,374 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
constant short FC_unary_op [[function_constant(FC_UNARY + 0)]];
|
||||
constant bool FC_unary_cnt[[function_constant(FC_UNARY + 1)]];
|
||||
|
||||
template <typename T0, typename T, typename TC>
|
||||
kernel void kernel_unary_impl(
|
||||
constant ggml_metal_kargs_unary & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort3 tpitg[[thread_position_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
#define FC_OP FC_unary_op
|
||||
#define FC_CNT FC_unary_cnt
|
||||
|
||||
device const T0 * src0_ptr;
|
||||
device T * dst_ptr;
|
||||
|
||||
int i0;
|
||||
|
||||
if (FC_CNT) {
|
||||
i0 = tgpig.x;
|
||||
|
||||
src0_ptr = (device const T0 *) (src0);
|
||||
dst_ptr = (device T *) (dst);
|
||||
} else {
|
||||
const int i03 = tgpig.z;
|
||||
const int i02 = tgpig.y;
|
||||
const int k0 = tgpig.x/args.ne01;
|
||||
const int i01 = tgpig.x - k0*args.ne01;
|
||||
|
||||
i0 = k0*ntg.x + tpitg.x;
|
||||
|
||||
src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
|
||||
dst_ptr = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1 );
|
||||
}
|
||||
|
||||
{
|
||||
//threadgroup_barrier(mem_flags::mem_none);
|
||||
|
||||
if (!FC_CNT) {
|
||||
if (i0 >= args.ne0) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
const TC x = (TC) src0_ptr[i0];
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SCALE) {
|
||||
dst_ptr[i0] = (T) (args.scale * x + args.bias);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_FILL) {
|
||||
dst_ptr[i0] = (T) args.val;
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_CLAMP) {
|
||||
dst_ptr[i0] = (T) clamp(x, args.min, args.max);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SQR) {
|
||||
dst_ptr[i0] = (T) (x * x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SQRT) {
|
||||
dst_ptr[i0] = (T) sqrt(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SIN) {
|
||||
dst_ptr[i0] = (T) sin(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_COS) {
|
||||
dst_ptr[i0] = (T) cos(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_LOG) {
|
||||
dst_ptr[i0] = (T) log(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_LEAKY_RELU) {
|
||||
dst_ptr[i0] = (T) (TC(x > 0)*x + TC(x <= 0)*(x * args.slope));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_TANH) {
|
||||
dst_ptr[i0] = (T) precise::tanh(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_RELU) {
|
||||
dst_ptr[i0] = (T) fmax(0, x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SIGMOID) {
|
||||
dst_ptr[i0] = (T) (1 / (1 + exp(-x)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_GELU) {
|
||||
dst_ptr[i0] = (T) (0.5*x*(1 + precise::tanh(SQRT_2_OVER_PI*x*(1 + GELU_COEF_A*x*x))));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_GELU_ERF) {
|
||||
dst_ptr[i0] = (T) (0.5*x*(1 + erf_approx(SQRT_2_INV*x)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_GELU_QUICK) {
|
||||
dst_ptr[i0] = (T) (x * (1/(1 + exp(GELU_QUICK_COEF*x))));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SILU) {
|
||||
dst_ptr[i0] = (T) (x / (1 + exp(-x)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_ELU) {
|
||||
dst_ptr[i0] = (T) elu_approx(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_NEG) {
|
||||
dst_ptr[i0] = (T) -x;
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_ABS) {
|
||||
dst_ptr[i0] = (T) fabs(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SGN) {
|
||||
dst_ptr[i0] = T(x > 0) - T(x < 0);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_STEP) {
|
||||
dst_ptr[i0] = T(x > 0);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_HARDSWISH) {
|
||||
dst_ptr[i0] = (T) (x * fmax(0, fmin(1, x/6 + 0.5)));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_HARDSIGMOID) {
|
||||
dst_ptr[i0] = (T) fmax(0, fmin(1, x/6 + 0.5));
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_EXP) {
|
||||
dst_ptr[i0] = (T) exp(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_SOFTPLUS) {
|
||||
dst_ptr[i0] = (T) select(log(1 + exp(x)), x, x > 20);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_EXPM1) {
|
||||
// TODO: precise implementation
|
||||
dst_ptr[i0] = (T) (exp(x) - 1);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_FLOOR) {
|
||||
dst_ptr[i0] = (T) floor(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_CEIL) {
|
||||
dst_ptr[i0] = (T) ceil(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_ROUND) {
|
||||
dst_ptr[i0] = (T) round(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_TRUNC) {
|
||||
dst_ptr[i0] = (T) trunc(x);
|
||||
}
|
||||
|
||||
if (FC_OP == OP_UNARY_NUM_XIELU) {
|
||||
const TC xi = x;
|
||||
const TC gate = TC(xi > TC(0.0f));
|
||||
const TC clamped = fmin(xi, TC(args.val));
|
||||
const TC y_pos = TC(args.scale) * xi * xi + TC(args.bias) * xi;
|
||||
const TC y_neg = (exp(clamped) - TC(1.0f) - xi) * TC(args.slope) + TC(args.bias) * xi;
|
||||
dst_ptr[i0] = (T) (gate * y_pos + (TC(1.0f) - gate) * y_neg);
|
||||
}
|
||||
}
|
||||
|
||||
#undef FC_OP
|
||||
#undef FC_CNT
|
||||
}
|
||||
|
||||
typedef decltype(kernel_unary_impl<float, float, float>) kernel_unary_t;
|
||||
|
||||
template [[host_name("kernel_unary_f32_f32")]] kernel kernel_unary_t kernel_unary_impl<float, float, float>;
|
||||
template [[host_name("kernel_unary_f32_f32_4")]] kernel kernel_unary_t kernel_unary_impl<float4, float4, float4>;
|
||||
template [[host_name("kernel_unary_f16_f16")]] kernel kernel_unary_t kernel_unary_impl<half, half, float>;
|
||||
template [[host_name("kernel_unary_f16_f16_4")]] kernel kernel_unary_t kernel_unary_impl<half4, half4, float4>;
|
||||
|
||||
kernel void kernel_silu_back_f32(
|
||||
constant ggml_metal_kargs_silu_back & args,
|
||||
device const float * dy,
|
||||
device const float * x,
|
||||
device float * dx,
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
if (gid >= args.ne) {
|
||||
return;
|
||||
}
|
||||
|
||||
const float s = 1.0f / (1.0f + exp(-x[gid]));
|
||||
dx[gid] = dy[gid] * s * (1.0f + x[gid] * (1.0f - s));
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_reglu(
|
||||
constant ggml_metal_kargs_glu & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
uint tpitg[[thread_position_in_threadgroup]],
|
||||
uint ntg[[threads_per_threadgroup]]) {
|
||||
device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
|
||||
device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
|
||||
device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
|
||||
|
||||
for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
|
||||
const float x0 = src0_row[i0];
|
||||
const float x1 = src1_row[i0];
|
||||
|
||||
dst_row[i0] = (T)(x0*x1*(x0 > 0.0f));
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_reglu<float>) kernel_reglu_t;
|
||||
|
||||
template [[host_name("kernel_reglu_f32")]] kernel kernel_reglu_t kernel_reglu<float>;
|
||||
template [[host_name("kernel_reglu_f16")]] kernel kernel_reglu_t kernel_reglu<half>;
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_geglu(
|
||||
constant ggml_metal_kargs_glu & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
uint tpitg[[thread_position_in_threadgroup]],
|
||||
uint ntg[[threads_per_threadgroup]]) {
|
||||
device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
|
||||
device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
|
||||
device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
|
||||
|
||||
for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
|
||||
const float x0 = src0_row[i0];
|
||||
const float x1 = src1_row[i0];
|
||||
|
||||
const float gelu = 0.5f*x0*(1.0f + precise::tanh(SQRT_2_OVER_PI*x0*(1.0f + GELU_COEF_A*x0*x0)));
|
||||
|
||||
dst_row[i0] = (T)(gelu*x1);
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_geglu<float>) kernel_geglu_t;
|
||||
|
||||
template [[host_name("kernel_geglu_f32")]] kernel kernel_geglu_t kernel_geglu<float>;
|
||||
template [[host_name("kernel_geglu_f16")]] kernel kernel_geglu_t kernel_geglu<half>;
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_swiglu(
|
||||
constant ggml_metal_kargs_glu & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
uint tpitg[[thread_position_in_threadgroup]],
|
||||
uint ntg[[threads_per_threadgroup]]) {
|
||||
device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
|
||||
device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
|
||||
device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
|
||||
|
||||
for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
|
||||
const float x0 = src0_row[i0];
|
||||
const float x1 = src1_row[i0];
|
||||
|
||||
const float silu = x0 / (1.0f + exp(-x0));
|
||||
|
||||
dst_row[i0] = (T)(silu*x1);
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_swiglu<float>) kernel_swiglu_t;
|
||||
|
||||
template [[host_name("kernel_swiglu_f32")]] kernel kernel_swiglu_t kernel_swiglu<float>;
|
||||
template [[host_name("kernel_swiglu_f16")]] kernel kernel_swiglu_t kernel_swiglu<half>;
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_swiglu_oai(
|
||||
constant ggml_metal_kargs_glu & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
uint tpitg[[thread_position_in_threadgroup]],
|
||||
uint ntg[[threads_per_threadgroup]]) {
|
||||
device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
|
||||
device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
|
||||
device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
|
||||
|
||||
for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
|
||||
float x0 = src0_row[i0];
|
||||
float x1 = src1_row[i0];
|
||||
|
||||
x0 = min(x0, args.limit);
|
||||
x1 = max(min(x1, args.limit), -args.limit);
|
||||
|
||||
float out_glu = x0 / (1.0f + exp(-x0 * args.alpha));
|
||||
out_glu = out_glu * (1.0f + x1);
|
||||
|
||||
dst_row[i0] = (T)out_glu;
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_swiglu_oai<float>) kernel_swiglu_oai_t;
|
||||
|
||||
template [[host_name("kernel_swiglu_oai_f32")]] kernel kernel_swiglu_oai_t kernel_swiglu_oai<float>;
|
||||
template [[host_name("kernel_swiglu_oai_f16")]] kernel kernel_swiglu_oai_t kernel_swiglu_oai<half>;
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_geglu_erf(
|
||||
constant ggml_metal_kargs_glu & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
uint tpitg[[thread_position_in_threadgroup]],
|
||||
uint ntg[[threads_per_threadgroup]]) {
|
||||
device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
|
||||
device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
|
||||
device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
|
||||
|
||||
for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
|
||||
const float x0 = src0_row[i0];
|
||||
const float x1 = src1_row[i0];
|
||||
|
||||
const float gelu_erf = 0.5f*x0*(1.0f+erf_approx<float>(x0*SQRT_2_INV));
|
||||
|
||||
dst_row[i0] = (T)(gelu_erf*x1);
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_geglu_erf<float>) kernel_geglu_erf_t;
|
||||
|
||||
template [[host_name("kernel_geglu_erf_f32")]] kernel kernel_geglu_erf_t kernel_geglu_erf<float>;
|
||||
template [[host_name("kernel_geglu_erf_f16")]] kernel kernel_geglu_erf_t kernel_geglu_erf<half>;
|
||||
|
||||
template<typename T>
|
||||
kernel void kernel_geglu_quick(
|
||||
constant ggml_metal_kargs_glu & args,
|
||||
device const char * src0,
|
||||
device const char * src1,
|
||||
device char * dst,
|
||||
uint tgpig[[threadgroup_position_in_grid]],
|
||||
uint tpitg[[thread_position_in_threadgroup]],
|
||||
uint ntg[[threads_per_threadgroup]]) {
|
||||
device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
|
||||
device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
|
||||
device T * dst_row = (device T *) ((device char *) dst + tgpig*args.nb1);
|
||||
|
||||
for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
|
||||
const float x0 = src0_row[i0];
|
||||
const float x1 = src1_row[i0];
|
||||
|
||||
const float gelu_quick = x0*(1.0f/(1.0f+exp(GELU_QUICK_COEF*x0)));
|
||||
|
||||
dst_row[i0] = (T)(gelu_quick*x1);
|
||||
}
|
||||
}
|
||||
|
||||
typedef decltype(kernel_geglu_quick<float>) kernel_geglu_quick_t;
|
||||
|
||||
template [[host_name("kernel_geglu_quick_f32")]] kernel kernel_geglu_quick_t kernel_geglu_quick<float>;
|
||||
template [[host_name("kernel_geglu_quick_f16")]] kernel kernel_geglu_quick_t kernel_geglu_quick<half>;
|
||||
@@ -1,179 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
constant bool FC_upscale_aa [[function_constant(FC_UPSCALE + 0)]];
|
||||
|
||||
kernel void kernel_upscale_nearest_f32(
|
||||
constant ggml_metal_kargs_upscale & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int64_t i3 = tgpig.z;
|
||||
const int64_t i2 = tgpig.y;
|
||||
const int64_t i1 = tgpig.x;
|
||||
|
||||
const int64_t i03 = i3/args.sf3;
|
||||
const int64_t i02 = i2/args.sf2;
|
||||
const int64_t i01 = i1/args.sf1;
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const int64_t i00 = i0/args.sf0;
|
||||
|
||||
device const float * src0_ptr = (device const float *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
|
||||
device float * dst_ptr = (device float *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
|
||||
|
||||
dst_ptr[0] = src0_ptr[0];
|
||||
}
|
||||
}
|
||||
|
||||
static inline float bilinear_tri(float x) {
|
||||
return MAX(0.0f, 1.0f - fabs(x));
|
||||
}
|
||||
|
||||
kernel void kernel_upscale_bilinear_f32(
|
||||
constant ggml_metal_kargs_upscale & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int64_t i3 = tgpig.z;
|
||||
const int64_t i2 = tgpig.y;
|
||||
const int64_t i1 = tgpig.x;
|
||||
|
||||
const int64_t i03 = i3 / args.sf3;
|
||||
const int64_t i02 = i2 / args.sf2;
|
||||
|
||||
const float f01 = ((float)i1 + args.poffs) / args.sf1 - args.poffs;
|
||||
const int64_t i01 = MAX(0, MIN(args.ne01 - 1, (int64_t)floor(f01)));
|
||||
const int64_t i01p = MAX(0, MIN(args.ne01 - 1, i01 + 1));
|
||||
const float fd1 = MAX(0.0f, MIN(1.0f, f01 - (float)i01));
|
||||
|
||||
src0 += i03*args.nb03 + i02*args.nb02;
|
||||
|
||||
device float * dst_ptr = (device float *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1);
|
||||
|
||||
if (FC_upscale_aa) {
|
||||
const float support0 = MAX(1.0f, 1.0f / args.sf0);
|
||||
const float invscale0 = 1.0f / support0;
|
||||
const float support1 = MAX(1.0f, 1.0f / args.sf1);
|
||||
const float invscale1 = 1.0f / support1;
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const float f00 = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
|
||||
|
||||
int64_t x_min = MAX((int64_t)0, (int64_t)floor(f00 - support0 + args.poffs));
|
||||
int64_t x_max = MIN(args.ne00, (int64_t)ceil (f00 + support0 + args.poffs));
|
||||
|
||||
int64_t y_min = MAX((int64_t)0, (int64_t)floor(f01 - support1 + args.poffs));
|
||||
int64_t y_max = MIN(args.ne01, (int64_t)ceil (f01 + support1 + args.poffs));
|
||||
|
||||
float sum = 0.0f;
|
||||
float wsum = 0.0f;
|
||||
|
||||
for (int64_t sy = y_min; sy < y_max; ++sy) {
|
||||
const float wy = MAX(0.0f, 1.0f - fabs((float)sy - f01) * invscale1);
|
||||
for (int64_t sx = x_min; sx < x_max; ++sx) {
|
||||
const float wx = MAX(0.0f, 1.0f - fabs((float)sx - f00) * invscale0);
|
||||
const float w = wx * wy;
|
||||
device const float * src_ptr = (device const float *)(src0 + sy*args.nb01 + sx*args.nb00);
|
||||
sum += (*src_ptr) * w;
|
||||
wsum += w;
|
||||
}
|
||||
}
|
||||
|
||||
const float v = (wsum > 0.0f) ? (sum / wsum) : 0.0f;
|
||||
dst_ptr[i0] = v;
|
||||
}
|
||||
} else {
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const float f00 = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
|
||||
const int64_t i00 = MAX(0, MIN(args.ne00 - 1, (int64_t)floor(f00)));
|
||||
const int64_t i00p = MAX(0, MIN(args.ne00 - 1, i00 + 1));
|
||||
const float fd0 = MAX(0.0f, MIN(1.0f, f00 - (float)i00));
|
||||
|
||||
device const float * src00 = (device const float *)(src0 + i01*args.nb01 + i00*args.nb00);
|
||||
device const float * src10 = (device const float *)(src0 + i01*args.nb01 + i00p*args.nb00);
|
||||
device const float * src01 = (device const float *)(src0 + i01p*args.nb01 + i00*args.nb00);
|
||||
device const float * src11 = (device const float *)(src0 + i01p*args.nb01 + i00p*args.nb00);
|
||||
|
||||
const float v =
|
||||
(*src00) * (1.0f - fd0) * (1.0f - fd1) +
|
||||
(*src10) * fd0 * (1.0f - fd1) +
|
||||
(*src01) * (1.0f - fd0) * fd1 +
|
||||
(*src11) * fd0 * fd1;
|
||||
|
||||
dst_ptr[i0] = v;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static inline float bicubic_weight1(float x) {
|
||||
const float a = -0.75f;
|
||||
return ((a + 2) * x - (a + 3)) * x * x + 1;
|
||||
}
|
||||
|
||||
static inline float bicubic_weight2(float x) {
|
||||
const float a = -0.75f;
|
||||
return ((a * x - 5 * a) * x + 8 * a) * x - 4 * a;
|
||||
}
|
||||
|
||||
kernel void kernel_upscale_bicubic_f32(
|
||||
constant ggml_metal_kargs_upscale & args,
|
||||
device const char * src0,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const int64_t i3 = tgpig.z;
|
||||
const int64_t i2 = tgpig.y;
|
||||
const int64_t i1 = tgpig.x;
|
||||
|
||||
const int64_t i03 = i3 / args.sf3;
|
||||
const int64_t i02 = i2 / args.sf2;
|
||||
|
||||
const float f01 = ((float)i1 + args.poffs) / args.sf1 - args.poffs;
|
||||
const int64_t i01 = (int64_t)floor(f01);
|
||||
const float fd1 = f01 - (float)i01;
|
||||
|
||||
const float w_y0 = bicubic_weight2(fd1 + 1.0f);
|
||||
const float w_y1 = bicubic_weight1(fd1);
|
||||
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;
|
||||
|
||||
device float * dst_ptr = (device float *)(dst + i3 * args.nb3 + i2 * args.nb2 + i1 * args.nb1);
|
||||
|
||||
for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
|
||||
const float f00 = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
|
||||
const int64_t i00 = (int64_t)floor(f00);
|
||||
const float fd0 = f00 - (float)i00;
|
||||
|
||||
const float w_x0 = bicubic_weight2(fd0 + 1.0f);
|
||||
const float w_x1 = bicubic_weight1(fd0);
|
||||
const float w_x2 = bicubic_weight1(1.0f - fd0);
|
||||
const float w_x3 = bicubic_weight2(2.0f - fd0);
|
||||
|
||||
float sum = 0.0f;
|
||||
|
||||
for (int dy = -1; dy <= 2; ++dy) {
|
||||
const int64_t iy = MAX(0, MIN(args.ne01 - 1, i01 + dy));
|
||||
const float wy = (dy == -1) ? w_y0 : (dy == 0) ? w_y1 : (dy == 1) ? w_y2 : w_y3;
|
||||
|
||||
for (int dx = -1; dx <= 2; ++dx) {
|
||||
const int64_t ix = MAX(0, MIN(args.ne00 - 1, i00 + dx));
|
||||
const float wx = (dx == -1) ? w_x0 : (dx == 0) ? w_x1 : (dx == 1) ? w_x2 : w_x3;
|
||||
|
||||
device const float * src_ptr = (device const float *)(src_slice + iy * args.nb01 + ix * args.nb00);
|
||||
sum += (*src_ptr) * wx * wy;
|
||||
}
|
||||
}
|
||||
|
||||
dst_ptr[i0] = sum;
|
||||
}
|
||||
}
|
||||
@@ -1,179 +0,0 @@
|
||||
#include "common.h"
|
||||
|
||||
kernel void kernel_rwkv_wkv6_f32(
|
||||
device const float * k,
|
||||
device const float * v,
|
||||
device const float * r,
|
||||
device const float * tf,
|
||||
device const float * td,
|
||||
device const float * state_in,
|
||||
device float * dst,
|
||||
constant uint & B,
|
||||
constant uint & T,
|
||||
constant uint & C,
|
||||
constant uint & H,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const uint head_size = 64; // TODO: support head_size = 128
|
||||
const uint batch_id = tgpig.x / H;
|
||||
const uint head_id = tgpig.x % H;
|
||||
const uint tid = tpitg.x;
|
||||
|
||||
if (batch_id >= B || head_id >= H) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint state_size = C * head_size;
|
||||
const uint n_seq_tokens = T / B;
|
||||
|
||||
threadgroup float _k[head_size];
|
||||
threadgroup float _r[head_size];
|
||||
threadgroup float _tf[head_size];
|
||||
threadgroup float _td[head_size];
|
||||
|
||||
float state[head_size];
|
||||
|
||||
for (uint i = 0; i < head_size; i++) {
|
||||
state[i] = state_in[batch_id * state_size + head_id * head_size * head_size
|
||||
+ i * head_size + tid];
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
_tf[tid] = tf[head_id * head_size + tid];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid;
|
||||
const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid;
|
||||
|
||||
for (uint t = start_t; t < end_t; t += C) {
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
_k[tid] = k[t];
|
||||
_r[tid] = r[t];
|
||||
_td[tid] = td[t];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
const float v_val = v[t];
|
||||
float y = 0.0;
|
||||
|
||||
for (uint j = 0; j < head_size; j += 4) {
|
||||
float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]);
|
||||
float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]);
|
||||
float4 tf_vec = float4(_tf[j], _tf[j+1], _tf[j+2], _tf[j+3]);
|
||||
float4 td_vec = float4(_td[j], _td[j+1], _td[j+2], _td[j+3]);
|
||||
float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
|
||||
|
||||
float4 kv = k_vec * v_val;
|
||||
|
||||
float4 temp = tf_vec * kv + s_vec;
|
||||
y += dot(r_vec, temp);
|
||||
|
||||
s_vec = s_vec * td_vec + kv;
|
||||
state[j] = s_vec[0];
|
||||
state[j+1] = s_vec[1];
|
||||
state[j+2] = s_vec[2];
|
||||
state[j+3] = s_vec[3];
|
||||
}
|
||||
|
||||
dst[t] = y;
|
||||
}
|
||||
|
||||
for (uint i = 0; i < head_size; i++) {
|
||||
dst[T * C + batch_id * state_size + head_id * head_size * head_size
|
||||
+ i * head_size + tid] = state[i];
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_rwkv_wkv7_f32(
|
||||
device const float * r,
|
||||
device const float * w,
|
||||
device const float * k,
|
||||
device const float * v,
|
||||
device const float * a,
|
||||
device const float * b,
|
||||
device const float * state_in,
|
||||
device float * dst,
|
||||
constant uint & B,
|
||||
constant uint & T,
|
||||
constant uint & C,
|
||||
constant uint & H,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
uint3 tpitg[[thread_position_in_threadgroup]],
|
||||
uint3 ntg[[threads_per_threadgroup]]) {
|
||||
|
||||
const uint head_size = 64; // TODO: support head_size = 128
|
||||
const uint batch_id = tgpig.x / H;
|
||||
const uint head_id = tgpig.x % H;
|
||||
const uint tid = tpitg.x;
|
||||
|
||||
if (batch_id >= B || head_id >= H) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint state_size = C * head_size;
|
||||
const uint n_seq_tokens = T / B;
|
||||
|
||||
threadgroup float _r[head_size];
|
||||
threadgroup float _w[head_size];
|
||||
threadgroup float _k[head_size];
|
||||
threadgroup float _a[head_size];
|
||||
threadgroup float _b[head_size];
|
||||
|
||||
float state[head_size];
|
||||
|
||||
for (uint i = 0; i < head_size; i++) {
|
||||
state[i] = state_in[batch_id * state_size + head_id * head_size * head_size
|
||||
+ tid * head_size + i];
|
||||
}
|
||||
|
||||
const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid;
|
||||
const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid;
|
||||
|
||||
for (uint t = start_t; t < end_t; t += C) {
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
_r[tid] = r[t];
|
||||
_w[tid] = w[t];
|
||||
_k[tid] = k[t];
|
||||
_a[tid] = a[t];
|
||||
_b[tid] = b[t];
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
const float v_val = v[t];
|
||||
float y = 0.0, sa = 0.0;
|
||||
|
||||
float4 sa_vec(0.0);
|
||||
|
||||
for (uint j = 0; j < head_size; j += 4) {
|
||||
float4 a_vec = float4(_a[j], _a[j+1], _a[j+2], _a[j+3]);
|
||||
float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
|
||||
sa_vec += a_vec * s_vec;
|
||||
}
|
||||
sa = sa_vec[0] + sa_vec[1] + sa_vec[2] + sa_vec[3];
|
||||
|
||||
for (uint j = 0; j < head_size; j += 4) {
|
||||
float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]);
|
||||
float4 w_vec = float4(_w[j], _w[j+1], _w[j+2], _w[j+3]);
|
||||
float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]);
|
||||
float4 b_vec = float4(_b[j], _b[j+1], _b[j+2], _b[j+3]);
|
||||
float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
|
||||
|
||||
float4 kv = k_vec * v_val;
|
||||
|
||||
s_vec = s_vec * w_vec + kv + sa * b_vec;
|
||||
y += dot(s_vec, r_vec);
|
||||
|
||||
state[j] = s_vec[0];
|
||||
state[j+1] = s_vec[1];
|
||||
state[j+2] = s_vec[2];
|
||||
state[j+3] = s_vec[3];
|
||||
}
|
||||
|
||||
dst[t] = y;
|
||||
}
|
||||
|
||||
for (uint i = 0; i < head_size; i++) {
|
||||
dst[T * C + batch_id * state_size + head_id * head_size * head_size
|
||||
+ tid * head_size + i] = state[i];
|
||||
}
|
||||
}
|
||||
@@ -323,6 +323,7 @@ class Keys:
|
||||
PROJECTOR_TYPE = "clip.projector_type"
|
||||
HAS_VISION_ENCODER = "clip.has_vision_encoder"
|
||||
HAS_AUDIO_ENCODER = "clip.has_audio_encoder"
|
||||
HAS_GEN_AUDIO_ENCODER = "clip.has_gen_audio_encoder"
|
||||
HAS_LLAVA_PROJECTOR = "clip.has_llava_projector"
|
||||
|
||||
class ClipVision:
|
||||
@@ -397,6 +398,18 @@ class Keys:
|
||||
DOWNSAMPLE_RATE = "clip.audio.projector.downsample_rate"
|
||||
HEAD_COUNT = "clip.audio.projector.head_count"
|
||||
|
||||
class ClipGenAudio:
|
||||
PROJECTOR_TYPE = "clip.gen.audio.projector_type" # for mixed modality models
|
||||
EMBEDDING_LENGTH = "clip.gen.audio.embedding_length"
|
||||
FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length"
|
||||
BLOCK_COUNT = "clip.gen.audio.block_count"
|
||||
PROJECTION_DIM = "clip.gen.audio.projection_dim"
|
||||
|
||||
class Attention:
|
||||
HEAD_COUNT = "clip.gen.audio.attention.head_count"
|
||||
HEAD_COUNT_KV = "clip.gen.audio.attention.head_count_kv"
|
||||
LAYERNORM_EPS = "clip.gen.audio.attention.layer_norm_epsilon"
|
||||
|
||||
class Diffusion:
|
||||
SHIFT_LOGITS = "diffusion.shift_logits"
|
||||
|
||||
@@ -558,6 +571,7 @@ class MODEL_ARCH(IntEnum):
|
||||
TALKIE = auto()
|
||||
MELLUM = auto()
|
||||
NANBEIGE = auto()
|
||||
QWEN3TTS = auto()
|
||||
|
||||
|
||||
class VISION_PROJECTOR_TYPE(IntEnum):
|
||||
@@ -958,6 +972,65 @@ class MODEL_TENSOR(IntEnum):
|
||||
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_ENC_CONV_RES2 = auto() # qwen3tts
|
||||
A_ENC_SE_CONV1 = auto() # qwen3tts
|
||||
A_ENC_SE_CONV2 = auto() # qwen3tts
|
||||
A_ENC_ASP_ATTN = auto() # qwen3tts
|
||||
A_ENC_ASP_TDNN = auto() # qwen3tts
|
||||
# qwen3tts code_predictor: predicts the remaining RVQ codebooks
|
||||
A_GEN_CODE_PROJ_IN = auto() # small_to_mtp_projection
|
||||
A_GEN_CODE_EMBD = auto() # per-codebook embedding table, merged 3D [n_codebooks, vocab, dim]
|
||||
A_GEN_CODE_HEAD = auto() # per-codebook output head, merged 3D [n_codebooks, vocab, dim]
|
||||
A_GEN_CODE_OUT_EMBD = auto() # codebook-0 embedding, re-fed into the talker backbone (talker.model.codec_embedding)
|
||||
A_GEN_CODE_ATTN_NORM = auto()
|
||||
A_GEN_CODE_ATTN_Q = auto()
|
||||
A_GEN_CODE_ATTN_Q_NORM = auto()
|
||||
A_GEN_CODE_ATTN_K = auto()
|
||||
A_GEN_CODE_ATTN_K_NORM = auto()
|
||||
A_GEN_CODE_ATTN_V = auto()
|
||||
A_GEN_CODE_ATTN_OUT = auto()
|
||||
A_GEN_CODE_FFN_NORM = auto()
|
||||
A_GEN_CODE_FFN_GATE = auto()
|
||||
A_GEN_CODE_FFN_UP = auto()
|
||||
A_GEN_CODE_FFN_DOWN = auto()
|
||||
A_GEN_CODE_OUTPUT_NORM = auto()
|
||||
# qwen3tts code2wav: RVQ codes -> raw PCM
|
||||
A_GEN_WAV_QUANT_FIRST_IN = auto() # semantic RVQ, in_proj (1x1 conv, loaded as 2D)
|
||||
A_GEN_WAV_QUANT_FIRST_OUT = auto() # semantic RVQ, out_proj
|
||||
A_GEN_WAV_QUANT_FIRST_CB = auto() # semantic RVQ codebook (1 layer), folded from embedding_sum/cluster_usage
|
||||
A_GEN_WAV_QUANT_REST_IN = auto() # acoustic RVQ, in_proj
|
||||
A_GEN_WAV_QUANT_REST_OUT = auto() # acoustic RVQ, out_proj
|
||||
A_GEN_WAV_QUANT_REST_CB = auto() # acoustic RVQ codebooks, merged 3D [15, vocab, dim]
|
||||
A_GEN_WAV_PRE_CONV = auto()
|
||||
A_GEN_WAV_TFM_IN_PROJ = auto()
|
||||
A_GEN_WAV_TFM_OUT_PROJ = auto()
|
||||
A_GEN_WAV_TFM_OUTPUT_NORM = auto()
|
||||
A_GEN_WAV_TFM_ATTN_NORM = auto()
|
||||
A_GEN_WAV_TFM_ATTN_Q = auto()
|
||||
A_GEN_WAV_TFM_ATTN_K = auto()
|
||||
A_GEN_WAV_TFM_ATTN_V = auto()
|
||||
A_GEN_WAV_TFM_ATTN_OUT = auto()
|
||||
A_GEN_WAV_TFM_ATTN_SCALE = auto() # layer scale (gamma) on the attn output
|
||||
A_GEN_WAV_TFM_FFN_NORM = auto()
|
||||
A_GEN_WAV_TFM_FFN_GATE = auto()
|
||||
A_GEN_WAV_TFM_FFN_UP = auto()
|
||||
A_GEN_WAV_TFM_FFN_DOWN = auto()
|
||||
A_GEN_WAV_TFM_FFN_SCALE = auto() # layer scale (gamma) on the FFN output
|
||||
A_GEN_WAV_UP_CONV = auto() # causal ConvTranspose1d, 2x upsample
|
||||
A_GEN_WAV_UP_DWCONV = auto() # ConvNeXt depthwise conv
|
||||
A_GEN_WAV_UP_NORM = auto() # ConvNeXt LayerNorm
|
||||
A_GEN_WAV_UP_PW1 = auto() # ConvNeXt pointwise conv 1 (expand)
|
||||
A_GEN_WAV_UP_PW2 = auto() # ConvNeXt pointwise conv 2 (project)
|
||||
A_GEN_WAV_UP_GAMMA = auto() # ConvNeXt layer scale
|
||||
A_GEN_WAV_DAC_ENTRY = auto() # DAC conv_pre
|
||||
A_GEN_WAV_DAC_UP_SNAKE = auto() # DAC per-block SnakeBeta before the upsample conv
|
||||
A_GEN_WAV_DAC_UP_CONV = auto() # DAC per-block causal ConvTranspose1d
|
||||
A_GEN_WAV_DAC_RES_ACT1 = auto() # DAC residual unit, SnakeBeta before conv1
|
||||
A_GEN_WAV_DAC_RES_CONV1 = auto() # DAC residual unit, dilated causal conv
|
||||
A_GEN_WAV_DAC_RES_ACT2 = auto() # DAC residual unit, SnakeBeta before conv2
|
||||
A_GEN_WAV_DAC_RES_CONV2 = auto() # DAC residual unit, pointwise causal conv
|
||||
A_GEN_WAV_DAC_POST_SNAKE = auto() # DAC final SnakeBeta
|
||||
A_GEN_WAV_DAC_POST_CONV = auto() # DAC conv_post -> 1-channel PCM
|
||||
A_MMPROJ = auto()
|
||||
A_MMPROJ_FC = auto()
|
||||
A_MM_NORM_PRE = auto()
|
||||
@@ -1170,6 +1243,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
|
||||
MODEL_ARCH.TALKIE: "talkie",
|
||||
MODEL_ARCH.MELLUM: "mellum",
|
||||
MODEL_ARCH.NANBEIGE: "nanbeige",
|
||||
MODEL_ARCH.QWEN3TTS: "qwen3tts",
|
||||
}
|
||||
|
||||
VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = {
|
||||
@@ -1567,6 +1641,63 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
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_ENC_CONV_RES2: "a.blk.{bid}.res2.{xid}",
|
||||
MODEL_TENSOR.A_ENC_SE_CONV1: "a.blk.{bid}.se_conv1",
|
||||
MODEL_TENSOR.A_ENC_SE_CONV2: "a.blk.{bid}.se_conv2",
|
||||
MODEL_TENSOR.A_ENC_ASP_ATTN: "a.asp_attn",
|
||||
MODEL_TENSOR.A_ENC_ASP_TDNN: "a.asp_tdnn",
|
||||
MODEL_TENSOR.A_GEN_CODE_PROJ_IN: "a.gen.code.proj_in",
|
||||
MODEL_TENSOR.A_GEN_CODE_EMBD: "a.gen.code.embd",
|
||||
MODEL_TENSOR.A_GEN_CODE_HEAD: "a.gen.code.head",
|
||||
MODEL_TENSOR.A_GEN_CODE_OUT_EMBD: "a.gen.code.out_embd",
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_NORM: "a.gen.code.blk.{bid}.ln1", # reuses the generic clip.cpp block loader (TN_LN_1)
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_Q: "a.gen.code.blk.{bid}.attn_q",
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM: "a.gen.code.blk.{bid}.attn_q_norm",
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_K: "a.gen.code.blk.{bid}.attn_k",
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM: "a.gen.code.blk.{bid}.attn_k_norm",
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_V: "a.gen.code.blk.{bid}.attn_v",
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_OUT: "a.gen.code.blk.{bid}.attn_out",
|
||||
MODEL_TENSOR.A_GEN_CODE_FFN_NORM: "a.gen.code.blk.{bid}.ln2", # reuses the generic clip.cpp block loader (TN_LN_2)
|
||||
MODEL_TENSOR.A_GEN_CODE_FFN_GATE: "a.gen.code.blk.{bid}.ffn_gate",
|
||||
MODEL_TENSOR.A_GEN_CODE_FFN_UP: "a.gen.code.blk.{bid}.ffn_up",
|
||||
MODEL_TENSOR.A_GEN_CODE_FFN_DOWN: "a.gen.code.blk.{bid}.ffn_down",
|
||||
MODEL_TENSOR.A_GEN_CODE_OUTPUT_NORM: "a.gen.code.output_norm",
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_IN: "a.gen.wav.quant.first.in_proj",
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_OUT: "a.gen.wav.quant.first.out_proj",
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_CB: "a.gen.wav.quant.first.codebook",
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_REST_IN: "a.gen.wav.quant.rest.in_proj",
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_REST_OUT: "a.gen.wav.quant.rest.out_proj",
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_REST_CB: "a.gen.wav.quant.rest.codebook",
|
||||
MODEL_TENSOR.A_GEN_WAV_PRE_CONV: "a.gen.wav.pre_conv",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_IN_PROJ: "a.gen.wav.tfm.in_proj",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_OUT_PROJ: "a.gen.wav.tfm.out_proj",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_OUTPUT_NORM: "a.gen.wav.tfm.output_norm",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM: "a.gen.wav.tfm.blk.{bid}.ln1",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q: "a.gen.wav.tfm.blk.{bid}.attn_q",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K: "a.gen.wav.tfm.blk.{bid}.attn_k",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V: "a.gen.wav.tfm.blk.{bid}.attn_v",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT: "a.gen.wav.tfm.blk.{bid}.attn_out",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE: "a.gen.wav.tfm.blk.{bid}.ls1",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM: "a.gen.wav.tfm.blk.{bid}.ln2",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_FFN_GATE: "a.gen.wav.tfm.blk.{bid}.ffn_gate",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP: "a.gen.wav.tfm.blk.{bid}.ffn_up",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN: "a.gen.wav.tfm.blk.{bid}.ffn_down",
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE: "a.gen.wav.tfm.blk.{bid}.ls2",
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_CONV: "a.gen.wav.up.blk.{bid}.conv",
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_DWCONV: "a.gen.wav.up.blk.{bid}.dwconv",
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_NORM: "a.gen.wav.up.blk.{bid}.norm",
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_PW1: "a.gen.wav.up.blk.{bid}.pw1",
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_PW2: "a.gen.wav.up.blk.{bid}.pw2",
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_GAMMA: "a.gen.wav.up.blk.{bid}.gamma",
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_ENTRY: "a.gen.wav.dac.entry",
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_UP_SNAKE: "a.gen.wav.dac.blk.{bid}.snake",
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_UP_CONV: "a.gen.wav.dac.blk.{bid}.conv",
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT1: "a.gen.wav.dac.blk.{bid}.res.{xid}.act1",
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV1: "a.gen.wav.dac.blk.{bid}.res.{xid}.conv1",
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT2: "a.gen.wav.dac.blk.{bid}.res.{xid}.act2",
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2: "a.gen.wav.dac.blk.{bid}.res.{xid}.conv2",
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE: "a.gen.wav.dac.post_snake",
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV: "a.gen.wav.dac.post_conv",
|
||||
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",
|
||||
@@ -1821,6 +1952,63 @@ 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_RES2,
|
||||
MODEL_TENSOR.A_ENC_SE_CONV1,
|
||||
MODEL_TENSOR.A_ENC_SE_CONV2,
|
||||
MODEL_TENSOR.A_ENC_ASP_ATTN,
|
||||
MODEL_TENSOR.A_ENC_ASP_TDNN,
|
||||
MODEL_TENSOR.A_GEN_CODE_PROJ_IN,
|
||||
MODEL_TENSOR.A_GEN_CODE_EMBD,
|
||||
MODEL_TENSOR.A_GEN_CODE_HEAD,
|
||||
MODEL_TENSOR.A_GEN_CODE_OUT_EMBD,
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_NORM,
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_Q,
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM,
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_K,
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM,
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_V,
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_OUT,
|
||||
MODEL_TENSOR.A_GEN_CODE_FFN_NORM,
|
||||
MODEL_TENSOR.A_GEN_CODE_FFN_GATE,
|
||||
MODEL_TENSOR.A_GEN_CODE_FFN_UP,
|
||||
MODEL_TENSOR.A_GEN_CODE_FFN_DOWN,
|
||||
MODEL_TENSOR.A_GEN_CODE_OUTPUT_NORM,
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_IN,
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_OUT,
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_CB,
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_REST_IN,
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_REST_OUT,
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_REST_CB,
|
||||
MODEL_TENSOR.A_GEN_WAV_PRE_CONV,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_IN_PROJ,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_OUT_PROJ,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_OUTPUT_NORM,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_FFN_GATE,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN,
|
||||
MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE,
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_CONV,
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_DWCONV,
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_NORM,
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_PW1,
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_PW2,
|
||||
MODEL_TENSOR.A_GEN_WAV_UP_GAMMA,
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_ENTRY,
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_UP_SNAKE,
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_UP_CONV,
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT1,
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV1,
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_RES_ACT2,
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2,
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE,
|
||||
MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV,
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_MEAN,
|
||||
MODEL_TENSOR.A_ENC_CONV_NORM_VAR,
|
||||
MODEL_TENSOR.A_ENC_MEL_FILTERS,
|
||||
@@ -4648,6 +4836,22 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
],
|
||||
MODEL_ARCH.QWEN3TTS: [
|
||||
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,
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
],
|
||||
}
|
||||
|
||||
# tensors that will not be serialized
|
||||
@@ -4922,6 +5126,8 @@ class VisionProjectorType:
|
||||
GLM4V = "glm4v"
|
||||
YOUTUVL = "youtuvl"
|
||||
NEMOTRON_V2_VL = "nemotron_v2_vl"
|
||||
QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder
|
||||
QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor
|
||||
HUNYUANVL = "hunyuanvl"
|
||||
PARAKEET = "parakeet" # audio
|
||||
MINIMAXM3 = "minimax_m3"
|
||||
|
||||
@@ -280,6 +280,10 @@ class GGUFWriter:
|
||||
|
||||
self.kv_data[0][key] = GGUFValue(value=val, type=vtype, sub_type=sub_type)
|
||||
|
||||
def remove_key(self, key: str) -> None:
|
||||
for kv_data in self.kv_data:
|
||||
kv_data.pop(key, None)
|
||||
|
||||
def add_uint8(self, key: str, val: int) -> None:
|
||||
self.add_key_value(key,val, GGUFValueType.UINT8)
|
||||
|
||||
@@ -1144,7 +1148,11 @@ class GGUFWriter:
|
||||
def add_precompiled_charsmap(self, charsmap: bytes) -> None:
|
||||
self.add_array(Keys.Tokenizer.PRECOMPILED_CHARSMAP, charsmap)
|
||||
|
||||
def add_chat_template(self, value: str | Sequence[Mapping[str, str]]) -> None:
|
||||
def add_chat_template(self, value: str | Sequence[Mapping[str, str]] | None) -> None:
|
||||
if value is None:
|
||||
self.remove_key(Keys.Tokenizer.CHAT_TEMPLATE)
|
||||
return
|
||||
|
||||
if not isinstance(value, str):
|
||||
template_default = None
|
||||
template_names = set()
|
||||
@@ -1199,6 +1207,9 @@ class GGUFWriter:
|
||||
def add_clip_has_audio_encoder(self, value: bool) -> None:
|
||||
self.add_bool(Keys.Clip.HAS_AUDIO_ENCODER, value)
|
||||
|
||||
def add_clip_has_gen_audio_encoder(self, value: bool) -> None:
|
||||
self.add_bool(Keys.Clip.HAS_GEN_AUDIO_ENCODER, value)
|
||||
|
||||
def add_clip_projector_type(self, value: str) -> None:
|
||||
self.add_string(Keys.Clip.PROJECTOR_TYPE, value)
|
||||
|
||||
@@ -1401,6 +1412,32 @@ class GGUFWriter:
|
||||
def add_audio_projector_head_count(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipAudio.Projector.HEAD_COUNT, value)
|
||||
|
||||
# audio generation (mmproj)
|
||||
|
||||
def add_clip_gen_audio_projector_type(self, value: str) -> None:
|
||||
self.add_string(Keys.ClipGenAudio.PROJECTOR_TYPE, value)
|
||||
|
||||
def add_gen_audio_projection_dim(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipGenAudio.PROJECTION_DIM, value)
|
||||
|
||||
def add_gen_audio_embedding_length(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipGenAudio.EMBEDDING_LENGTH, value)
|
||||
|
||||
def add_gen_audio_feed_forward_length(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipGenAudio.FEED_FORWARD_LENGTH, value)
|
||||
|
||||
def add_gen_audio_block_count(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipGenAudio.BLOCK_COUNT, value)
|
||||
|
||||
def add_gen_audio_head_count(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipGenAudio.Attention.HEAD_COUNT, value)
|
||||
|
||||
def add_gen_audio_head_count_kv(self, value: int) -> None:
|
||||
self.add_uint32(Keys.ClipGenAudio.Attention.HEAD_COUNT_KV, value)
|
||||
|
||||
def add_gen_audio_attention_layernorm_eps(self, value: float) -> None:
|
||||
self.add_float32(Keys.ClipGenAudio.Attention.LAYERNORM_EPS, value)
|
||||
|
||||
def add_xielu_alpha_p(self, values: Sequence[float]):
|
||||
self.add_array(Keys.xIELU.ALPHA_P, values)
|
||||
|
||||
|
||||
@@ -2109,6 +2109,7 @@ class TensorNameMap:
|
||||
"conformer.subsample_conv_projection.layer{bid}.conv", # gemma4
|
||||
"sound_encoder.encoder.subsampling.layers.{bid}", # parakeet
|
||||
"encoder.conv{bid}", # mimo-audio-tokenizer
|
||||
"speaker_encoder.blocks.{bid}.conv", # qwen3tts speaker encoder (only bid=0, the stem TDNN)
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV1D_NORM: (
|
||||
@@ -2126,6 +2127,7 @@ class TensorNameMap:
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_OUT: (
|
||||
"audio_tower.conv_out", # qwen3omni
|
||||
"speaker_encoder.mfa.conv", # qwen3tts speaker encoder: multi-layer feature aggregation
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_PRE_NORM: (),
|
||||
@@ -2336,7 +2338,8 @@ class TensorNameMap:
|
||||
MODEL_TENSOR.A_MMPROJ_FC: (
|
||||
"audio.multi_modal_projector.linear", # qwen2audio
|
||||
"audio_tower.proj", # qwen2omni
|
||||
"model.audio_tower.output_proj" # gemma4
|
||||
"model.audio_tower.output_proj", # gemma4
|
||||
"speaker_encoder.fc", # qwen3tts speaker encoder: final speaker embedding projection
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_MM_NORM_PRE: (
|
||||
@@ -2411,6 +2414,7 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.lconv1d.linear_start", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.conv.pointwise_conv1", # parakeet
|
||||
"encoder.layers.{bid}.conv.up_conv", # granite_speech
|
||||
"speaker_encoder.blocks.{bid}.tdnn1.conv", # qwen3tts speaker encoder
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_CONV_PW2: (
|
||||
@@ -2418,6 +2422,23 @@ class TensorNameMap:
|
||||
"conformer.layers.{bid}.lconv1d.linear_end", # gemma3n
|
||||
"sound_encoder.encoder.layers.{bid}.conv.pointwise_conv2", # parakeet
|
||||
"encoder.layers.{bid}.conv.down_conv", # granite_speech
|
||||
"speaker_encoder.blocks.{bid}.tdnn2.conv", # qwen3tts speaker encoder
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_SE_CONV1: (
|
||||
"speaker_encoder.blocks.{bid}.se_block.conv1", # qwen3tts
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_SE_CONV2: (
|
||||
"speaker_encoder.blocks.{bid}.se_block.conv2", # qwen3tts
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_ASP_ATTN: (
|
||||
"speaker_encoder.asp.conv", # qwen3tts
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_ASP_TDNN: (
|
||||
"speaker_encoder.asp.tdnn.conv", # qwen3tts
|
||||
),
|
||||
|
||||
MODEL_TENSOR.A_ENC_NORM_CONV: (
|
||||
|
||||
@@ -119,6 +119,7 @@ Public API changes carry a higher bar than internal ones (`CONTRIBUTING.md`). Re
|
||||
- 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.
|
||||
- For audio generation models, see `tools/mtmd/README-dev.md`
|
||||
|
||||
## General (always)
|
||||
|
||||
|
||||
@@ -144,6 +144,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
|
||||
{ LLM_ARCH_TALKIE, "talkie" },
|
||||
{ LLM_ARCH_MELLUM, "mellum" },
|
||||
{ LLM_ARCH_NANBEIGE, "nanbeige" },
|
||||
{ LLM_ARCH_QWEN3TTS, "qwen3tts" },
|
||||
{ LLM_ARCH_UNKNOWN, "(unknown)" },
|
||||
};
|
||||
|
||||
@@ -1026,6 +1027,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
|
||||
case LLM_ARCH_MINIMAX_M3:
|
||||
case LLM_ARCH_MISTRAL4:
|
||||
case LLM_ARCH_KIMI_LINEAR:
|
||||
case LLM_ARCH_QWEN3TTS:
|
||||
return false;
|
||||
default:
|
||||
return true;
|
||||
|
||||
@@ -149,6 +149,7 @@ enum llm_arch {
|
||||
LLM_ARCH_MINIMAX_M3,
|
||||
LLM_ARCH_DFLASH,
|
||||
LLM_ARCH_NANBEIGE,
|
||||
LLM_ARCH_QWEN3TTS,
|
||||
LLM_ARCH_UNKNOWN,
|
||||
};
|
||||
|
||||
|
||||
@@ -124,3 +124,9 @@ LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx);
|
||||
LLAMA_API const int32_t * llama_model_target_layer_ids (const struct llama_model * model);
|
||||
// returns the number of extracted layers from target model
|
||||
LLAMA_API uint32_t llama_model_target_layer_ids_n(const struct llama_model * model);
|
||||
|
||||
// retrieves the whole token embedding matrix in F32 format (n_embd * n_vocab)
|
||||
// returns total number of elements or 0 on error
|
||||
// if out is nullptr, returns the number of tokens without writing to out
|
||||
// caller must allocate enough memory for out before calling
|
||||
LLAMA_API uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out);
|
||||
|
||||
@@ -112,6 +112,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
|
||||
return new llama_model_qwen3vl(params);
|
||||
case LLM_ARCH_QWEN3VLMOE:
|
||||
return new llama_model_qwen3vlmoe(params);
|
||||
case LLM_ARCH_QWEN3TTS:
|
||||
return new llama_model_qwen3tts(params);
|
||||
case LLM_ARCH_PHI2:
|
||||
return new llama_model_phi2(params);
|
||||
case LLM_ARCH_PHI3:
|
||||
@@ -2693,6 +2695,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
|
||||
case LLM_ARCH_QWEN3VLMOE:
|
||||
case LLM_ARCH_QWEN35:
|
||||
case LLM_ARCH_QWEN35MOE:
|
||||
case LLM_ARCH_QWEN3TTS:
|
||||
return LLAMA_ROPE_TYPE_IMROPE;
|
||||
|
||||
case LLM_ARCH_GLM4:
|
||||
@@ -2908,3 +2911,38 @@ const int32_t * llama_model_target_layer_ids(const struct llama_model * model) {
|
||||
uint32_t llama_model_target_layer_ids_n(const struct llama_model * model) {
|
||||
return (uint32_t) model->target_layer_ids.size();
|
||||
}
|
||||
|
||||
uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out) {
|
||||
if (model->vocab.n_tokens() == 0 || model->tok_embd == nullptr) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const ggml_tensor * tensor = model->tok_embd;
|
||||
const size_t nelements = ggml_nelements(tensor);
|
||||
GGML_ASSERT(nelements <= UINT32_MAX); // for the return type
|
||||
|
||||
if (out == nullptr) {
|
||||
return (uint32_t) nelements;
|
||||
}
|
||||
|
||||
if (tensor->type == GGML_TYPE_F32) {
|
||||
ggml_backend_tensor_get(tensor, out, 0, nelements * sizeof(float));
|
||||
return (uint32_t) nelements;
|
||||
}
|
||||
|
||||
std::vector<uint8_t> buf(ggml_nbytes(tensor));
|
||||
ggml_backend_tensor_get(tensor, buf.data(), 0, buf.size());
|
||||
|
||||
const ggml_type_traits * traits = ggml_get_type_traits(tensor->type);
|
||||
if (tensor->type == GGML_TYPE_F16) {
|
||||
ggml_fp16_to_fp32_row((const ggml_fp16_t *) buf.data(), out, nelements);
|
||||
} else if (tensor->type == GGML_TYPE_BF16) {
|
||||
ggml_bf16_to_fp32_row((const ggml_bf16_t *) buf.data(), out, nelements);
|
||||
} else if (ggml_is_quantized(tensor->type) && traits->to_float != nullptr) {
|
||||
traits->to_float(buf.data(), out, nelements);
|
||||
} else {
|
||||
GGML_ABORT("unsupported tensor type for dequantization: %s", ggml_type_name(tensor->type));
|
||||
}
|
||||
|
||||
return (uint32_t) nelements;
|
||||
}
|
||||
|
||||
@@ -125,7 +125,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
|
||||
layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0);
|
||||
layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0);
|
||||
layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0);
|
||||
layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank * o_groups}, 0);
|
||||
layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank, o_groups}, TENSOR_ALLOW_RESHAPE);
|
||||
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0);
|
||||
|
||||
layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
|
||||
|
||||
@@ -596,6 +596,11 @@ struct llama_model_qwen3vlmoe : public llama_model_base {
|
||||
};
|
||||
|
||||
|
||||
struct llama_model_qwen3tts : public llama_model_qwen3vl {
|
||||
llama_model_qwen3tts(const struct llama_model_params & params) : llama_model_qwen3vl(params) {}
|
||||
};
|
||||
|
||||
|
||||
struct llama_model_phi2 : public llama_model_base {
|
||||
llama_model_phi2(const struct llama_model_params & params) : llama_model_base(params) {}
|
||||
void load_arch_hparams(llama_model_loader & ml) override;
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
#include "models.h"
|
||||
|
||||
// llama_model_qwen3tts reuses llama_model_qwen3vl's hparams/tensors/graph logic
|
||||
+24
-1
@@ -16,11 +16,16 @@ void llama_model_qwen3vl::load_arch_hparams(llama_model_loader & ml) {
|
||||
void llama_model_qwen3vl::load_arch_tensors(llama_model_loader &) {
|
||||
LLAMA_LOAD_LOCALS;
|
||||
|
||||
int64_t n_vocab_out = n_vocab;
|
||||
if (arch == LLM_ARCH_QWEN3TTS) {
|
||||
n_vocab_out = 3072;
|
||||
}
|
||||
|
||||
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
|
||||
|
||||
// output
|
||||
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
|
||||
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
|
||||
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab_out}, TENSOR_NOT_REQUIRED);
|
||||
// if output is NULL, init from the input tok embed
|
||||
if (output == NULL) {
|
||||
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
|
||||
@@ -166,6 +171,24 @@ llama_model_qwen3vl::graph::graph(const llama_model & model, const llm_graph_par
|
||||
// lm_head
|
||||
cur = build_lora_mm(model.output, cur, model.output_s);
|
||||
|
||||
int64_t n_vocab_in = model.tok_embd->ne[1];
|
||||
int64_t n_vocab_out = model.output->ne[1];
|
||||
if (n_vocab_in > n_vocab_out) {
|
||||
// case: Qwen3TTS model with codec_head as output
|
||||
GGML_ASSERT(model.output_norm);
|
||||
int64_t pad = n_vocab_in - n_vocab_out;
|
||||
|
||||
// using this trick to get a scalar -inf tensor to pad the output
|
||||
ggml_tensor * neg_inf = ggml_scale_bias(ctx0,
|
||||
ggml_view_1d(ctx0, model.output_norm, 1, 0),
|
||||
0.0f, -INFINITY);
|
||||
neg_inf = ggml_repeat_4d(ctx0, neg_inf, pad, cur->ne[1], 1, 1);
|
||||
cur = ggml_concat(ctx0, neg_inf, cur, 0); // [padded .. n_vocab_out, n_stream]
|
||||
|
||||
} else if (n_vocab_in < n_vocab_out) {
|
||||
GGML_ABORT("invalid case");
|
||||
}
|
||||
|
||||
cb(cur, "result_output", -1);
|
||||
res->t_logits = cur;
|
||||
|
||||
|
||||
@@ -113,6 +113,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|
||||
n_layer = 3;
|
||||
} else if (arch == LLM_ARCH_CHAMELEON) {
|
||||
n_vocab = 10240;
|
||||
} else if (arch == LLM_ARCH_QWEN3TTS) {
|
||||
n_vocab = 4096; // must be >= the hard-coded codec head size (3072)
|
||||
}
|
||||
|
||||
const uint32_t n_embd_head = n_embd / n_head;
|
||||
|
||||
@@ -18,6 +18,8 @@ add_library(mtmd
|
||||
mtmd-image.cpp
|
||||
mtmd.h
|
||||
mtmd-helper.cpp
|
||||
mtmd-helper-gen.cpp
|
||||
mtmd-helper-common.h
|
||||
mtmd-helper.h
|
||||
clip.cpp
|
||||
clip.h
|
||||
@@ -52,6 +54,8 @@ add_library(mtmd
|
||||
models/mimovl.cpp
|
||||
models/qwen3a.cpp
|
||||
models/mimo-audio.cpp
|
||||
models/qwen3tts-spkenc.cpp
|
||||
models/qwen3tts-gen.cpp
|
||||
models/step3vl.cpp
|
||||
models/siglip.cpp
|
||||
models/whisper-enc.cpp
|
||||
|
||||
@@ -33,3 +33,52 @@ A typical pipeline of the core libmtmd is as follows:
|
||||
We provide a set of helper functions via `mtmd_helper` to make using libmtmd easier. The helper provides:
|
||||
- Image, audio and video file decoding (for example, decode raw JPEG into RGB bitmap)
|
||||
- Manage `llama_batch` and calls to `llama_decode`
|
||||
|
||||
## Audio generation support
|
||||
|
||||
Audio generation is added to mtmd in PR [#26254](https://github.com/ggml-org/llama.cpp/pull/26254)
|
||||
|
||||
Currently, we support the 3-stage pipeline below which should cover most TTS models:
|
||||
- Stage 1: Backbone / Semantic Stage: Backbone model accepts text prompt and reference voice as input
|
||||
- Stage 2: Acoustic Detail Generator: A model takes the hidden state from backbone and generate audio details (usually as audio codes or mel-spectrogram)
|
||||
- Stage 3: Waveform Reconstruction: Convert the semantic and acoustic data from previous stages to the final waveform
|
||||
|
||||
For example, Qwen3-TTS:
|
||||
- Reference voice is encoded using ECAPA-TDNN speaker encoder (`speaker_encoder`)
|
||||
- Text prompt and reference voice are processed via a backbone (`talker.model`)
|
||||
- A model converts sampled semantic token and hidden state from stage 2 into a list of 15 acoustic codes (`talker.code_predictor`)
|
||||
- 16 generated codes are converted into waveform (`code2wav`)
|
||||
|
||||
### API design constraints
|
||||
|
||||
Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system is designed to be flexible and reusable by new models.
|
||||
|
||||
`mtmd_gen_audio` is split into 2 main API:
|
||||
- Core API `mtmd.h`: handles main inference. Important: the API surface must be stateless; caller must handle state management and audio frame accumulation.
|
||||
- Helper API `mtmd-helper.h`: provides a model-agnostic stateful API. Usage example can be found in the `tools/tts` directory.
|
||||
|
||||
### Checklist for porting new audio generation models to mtmd
|
||||
|
||||
1. Establish a list of reusable and missing components from the current mtmd implementation.
|
||||
2. For GGUF conversion:
|
||||
- Backbone model should be converted to a normal text model (loadable via `libllama`)
|
||||
- If model used hard-coded embedding row ID, append them to token embeddings and assign token name for them (see `qwen3tts.py`)
|
||||
- If model have a specific output logits head for audio codes (usually semantic code), keep the head as-is and pad the logits at inference time (see `src/models/qwen3vl.cpp`)
|
||||
- Sidecar models (code2wav, bigvgan, etc) must live inside the mmproj GGUF (but can be in different `clip_context` if necessary)
|
||||
- Note: it should use `ggml_build_forward_select` to select graphs if multiple graphs living in the same context
|
||||
- Reuse existing GGUF metadata key name and tensor name whenever possible; think twice before adding extensive changes to GGUF writer. For example, Qwen3-TTS hard-code part of the hparams to `clip.cpp` as they won't likely to change.
|
||||
- For tensor naming:
|
||||
- Prefixed with `a.*` for tensors used by speaker encoder pipeline
|
||||
- Prefixed with `a.gen.*` for generation stages (code / mel-spectrogram / PCM generation)
|
||||
3. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this:
|
||||
- 10-20% changes is to add new backbone (text) model and conversion
|
||||
- 60% changes inside `mtmd-helper-gen.cpp`
|
||||
- 10% changes inside `libmtmd` and `clip.cpp` systems
|
||||
- The rest downstream code (CLI, server) should have no changes at all
|
||||
4. Update usage documentation in `tools/tts/README.md`
|
||||
|
||||
IMPORTANT: If your model needs changes that don't fit the existing infrastructure, **open an issue first for discussion**.
|
||||
|
||||
No-go checklist (these will get the PR rejected and require discussion before proceeding):
|
||||
- Violating the API design constraints stated above
|
||||
- Adding a new model-specific binary: the API and binary surface must stay model-agnostic
|
||||
|
||||
@@ -54,6 +54,9 @@ struct clip_graph {
|
||||
|
||||
clip_graph(clip_ctx * ctx, const clip_image_f32 & img);
|
||||
|
||||
// build sub-graph, reuse buf from parent
|
||||
clip_graph(const clip_graph & parent);
|
||||
|
||||
virtual ~clip_graph() = default;
|
||||
virtual ggml_cgraph * build() = 0;
|
||||
|
||||
|
||||
@@ -32,6 +32,7 @@
|
||||
#define KEY_PROJ_TYPE "clip.projector_type"
|
||||
#define KEY_HAS_AUDIO_ENC "clip.has_audio_encoder"
|
||||
#define KEY_HAS_VISION_ENC "clip.has_vision_encoder"
|
||||
#define KEY_HAS_GEN_AUDIO_ENC "clip.has_gen_audio_encoder"
|
||||
#define KEY_USE_GELU "clip.use_gelu"
|
||||
#define KEY_USE_SILU "clip.use_silu"
|
||||
|
||||
@@ -89,6 +90,8 @@
|
||||
#define KEY_A_ATTN_WINDOW_SIZE "clip.audio.window_size" // mimo-audio-tokenizer: sliding-window radius
|
||||
#define KEY_A_LOCAL_BLOCK_COUNT "clip.audio.local_block_count" // mimo-v2.5: input_local_transformer layer count
|
||||
#define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size
|
||||
// audio generation (gen-audio)-specific
|
||||
#define KEY_GEN_AUDIO_PROJ_TYPE "clip.gen.audio.projector_type" // for models with mixed modalities
|
||||
#define KEY_AUDIO_SUBSAMPLING_FACTOR "clip.audio.subsampling_factor"
|
||||
|
||||
//
|
||||
@@ -201,6 +204,48 @@
|
||||
#define TN_MM_A_LOCAL_LN2 "mm.a.local_blk.%d.ln2.%s"
|
||||
#define TN_MM_A_LOCAL_NORM "mm.a.local_norm.%s"
|
||||
|
||||
// qwen3tts speaker encoder (ECAPA-TDNN)
|
||||
#define TN_A_SE_CONV1 "a.blk.%d.se_conv1.%s"
|
||||
#define TN_A_SE_CONV2 "a.blk.%d.se_conv2.%s"
|
||||
#define TN_A_CONV_RES2 "a.blk.%d.res2.%d.%s"
|
||||
#define TN_A_ASP_ATTN "a.asp_attn.%s"
|
||||
#define TN_A_ASP_TDNN "a.asp_tdnn.%s"
|
||||
|
||||
// qwen3tts code_predictor
|
||||
#define TN_A_GEN_CODE_PROJ_IN "a.gen.code.proj_in.%s"
|
||||
#define TN_A_GEN_CODE_EMBD "a.gen.code.embd.%s"
|
||||
#define TN_A_GEN_CODE_HEAD "a.gen.code.head.%s"
|
||||
#define TN_A_GEN_CODE_OUT_EMBD "a.gen.code.out_embd.%s"
|
||||
#define TN_A_GEN_CODE_NORM "a.gen.code.output_norm.%s"
|
||||
|
||||
// qwen3tts code2wav (RVQ codes -> raw PCM)
|
||||
// pre_transformer layers use the generic TN_ATTN_*/TN_FFN_*/TN_LN_*/TN_LS_* macros, prefix "a.gen.wav.tfm"
|
||||
#define TN_A_GEN_WAV_QUANT_FIRST_IN "a.gen.wav.quant.first.in_proj.%s"
|
||||
#define TN_A_GEN_WAV_QUANT_FIRST_OUT "a.gen.wav.quant.first.out_proj.%s"
|
||||
#define TN_A_GEN_WAV_QUANT_FIRST_CB "a.gen.wav.quant.first.codebook.%s"
|
||||
#define TN_A_GEN_WAV_QUANT_REST_IN "a.gen.wav.quant.rest.in_proj.%s"
|
||||
#define TN_A_GEN_WAV_QUANT_REST_OUT "a.gen.wav.quant.rest.out_proj.%s"
|
||||
#define TN_A_GEN_WAV_QUANT_REST_CB "a.gen.wav.quant.rest.codebook.%s"
|
||||
#define TN_A_GEN_WAV_PRE_CONV "a.gen.wav.pre_conv.%s"
|
||||
#define TN_A_GEN_WAV_TFM_IN_PROJ "a.gen.wav.tfm.in_proj.%s"
|
||||
#define TN_A_GEN_WAV_TFM_OUT_PROJ "a.gen.wav.tfm.out_proj.%s"
|
||||
#define TN_A_GEN_WAV_TFM_OUT_NORM "a.gen.wav.tfm.output_norm.%s"
|
||||
#define TN_A_GEN_WAV_UP_CONV "a.gen.wav.up.blk.%d.conv.%s"
|
||||
#define TN_A_GEN_WAV_UP_DWCONV "a.gen.wav.up.blk.%d.dwconv.%s"
|
||||
#define TN_A_GEN_WAV_UP_NORM "a.gen.wav.up.blk.%d.norm.%s"
|
||||
#define TN_A_GEN_WAV_UP_PW1 "a.gen.wav.up.blk.%d.pw1.%s"
|
||||
#define TN_A_GEN_WAV_UP_PW2 "a.gen.wav.up.blk.%d.pw2.%s"
|
||||
#define TN_A_GEN_WAV_UP_GAMMA "a.gen.wav.up.blk.%d.gamma"
|
||||
#define TN_A_GEN_WAV_DAC_ENTRY "a.gen.wav.dac.entry.%s"
|
||||
#define TN_A_GEN_WAV_DAC_SNAKE "a.gen.wav.dac.blk.%d.snake.%s"
|
||||
#define TN_A_GEN_WAV_DAC_CONV "a.gen.wav.dac.blk.%d.conv.%s"
|
||||
#define TN_A_GEN_WAV_DAC_RES_ACT1 "a.gen.wav.dac.blk.%d.res.%d.act1.%s"
|
||||
#define TN_A_GEN_WAV_DAC_RES_CONV1 "a.gen.wav.dac.blk.%d.res.%d.conv1.%s"
|
||||
#define TN_A_GEN_WAV_DAC_RES_ACT2 "a.gen.wav.dac.blk.%d.res.%d.act2.%s"
|
||||
#define TN_A_GEN_WAV_DAC_RES_CONV2 "a.gen.wav.dac.blk.%d.res.%d.conv2.%s"
|
||||
#define TN_A_GEN_WAV_DAC_POST_SNAKE "a.gen.wav.dac.post_snake.%s"
|
||||
#define TN_A_GEN_WAV_DAC_POST_CONV "a.gen.wav.dac.post_conv.%s"
|
||||
|
||||
// cogvlm
|
||||
#define TN_MM_POST_FC_NORM "mm.post_fc_norm.%s"
|
||||
#define TN_MM_H_TO_4H "mm.up.%s"
|
||||
@@ -408,6 +453,8 @@ enum projector_type {
|
||||
PROJECTOR_TYPE_MINIMAX_M3,
|
||||
PROJECTOR_TYPE_GRANITE4_VISION,
|
||||
PROJECTOR_TYPE_MIMO_AUDIO,
|
||||
PROJECTOR_TYPE_QWEN3TTS_SPKENC,
|
||||
PROJECTOR_TYPE_QWEN3TTS_GEN,
|
||||
PROJECTOR_TYPE_UNKNOWN,
|
||||
};
|
||||
|
||||
@@ -465,6 +512,8 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
|
||||
{ PROJECTOR_TYPE_GRANITE4_VISION, "granite4_vision"},
|
||||
{ PROJECTOR_TYPE_MIMO_AUDIO, "mimo_audio"},
|
||||
{ PROJECTOR_TYPE_PARAKEET, "parakeet"},
|
||||
{ PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"},
|
||||
{ PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"},
|
||||
};
|
||||
|
||||
static projector_type clip_projector_type_from_string(const std::string & str) {
|
||||
|
||||
@@ -136,6 +136,19 @@ struct clip_hparams {
|
||||
int32_t rvq_num_quantizers = 0;
|
||||
std::vector<int32_t> rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17)
|
||||
|
||||
// qwen3tts code2wav
|
||||
int32_t wav_tfm_n_layer = 0;
|
||||
int32_t wav_tfm_n_embd = 0;
|
||||
int32_t wav_tfm_n_ff = 0;
|
||||
int32_t wav_tfm_n_head = 0;
|
||||
int32_t wav_tfm_n_head_kv = 0;
|
||||
float wav_tfm_eps = 1e-5f;
|
||||
float wav_tfm_rope_theta = 10000.0f;
|
||||
int32_t wav_upsample_n_block = 0;
|
||||
int32_t wav_dac_n_block = 0;
|
||||
int32_t wav_dac_n_res = 0;
|
||||
int32_t wav_tfm_swa = 0; // pre_transformer's KV cache size, in frames
|
||||
|
||||
// mimo-v2.5: LLM-side connector (input_local_transformer)
|
||||
int32_t audio_local_n_layer = 0;
|
||||
int32_t audio_local_group_size = 0;
|
||||
@@ -286,6 +299,14 @@ struct clip_layer {
|
||||
ggml_tensor * cross_attn_norm_w = nullptr;
|
||||
ggml_tensor * cross_attn_norm_b = nullptr;
|
||||
|
||||
// qwen3tts speaker encoder: SE-Res2Net block, tdnn1/tdnn2 reuse conv_pw1_w/b and conv_pw2_w/b above
|
||||
ggml_tensor * se_conv1_w = nullptr;
|
||||
ggml_tensor * se_conv1_b = nullptr;
|
||||
ggml_tensor * se_conv2_w = nullptr;
|
||||
ggml_tensor * se_conv2_b = nullptr;
|
||||
std::vector<ggml_tensor *> res2_conv_w; // Res2Net hierarchical branches
|
||||
std::vector<ggml_tensor *> res2_conv_b;
|
||||
|
||||
bool has_deepstack() const {
|
||||
return deepstack_fc1_w != nullptr;
|
||||
}
|
||||
@@ -365,6 +386,73 @@ struct qf_block {
|
||||
std::vector<clip_layer> qf_proj_layers;
|
||||
};
|
||||
|
||||
// qwen3tts code2wav: RVQ codes -> raw PCM
|
||||
struct clip_code2wav {
|
||||
// "upsample" stage: one ConvNeXt block plus the causal ConvTranspose1d before it
|
||||
struct upsample_block {
|
||||
ggml_tensor * conv_w = nullptr; // causal ConvTranspose1d, 2x
|
||||
ggml_tensor * conv_b = nullptr;
|
||||
ggml_tensor * dwconv_w = nullptr; // depthwise causal conv, k=7
|
||||
ggml_tensor * dwconv_b = nullptr;
|
||||
ggml_tensor * norm_w = nullptr; // LayerNorm
|
||||
ggml_tensor * norm_b = nullptr;
|
||||
ggml_tensor * pw1_w = nullptr; // pointwise expand
|
||||
ggml_tensor * pw1_b = nullptr;
|
||||
ggml_tensor * pw2_w = nullptr; // pointwise project
|
||||
ggml_tensor * pw2_b = nullptr;
|
||||
ggml_tensor * gamma = nullptr; // layer scale
|
||||
};
|
||||
|
||||
// one DAC residual unit: SnakeBeta -> dilated causal conv -> SnakeBeta -> pointwise causal conv
|
||||
struct dac_res {
|
||||
ggml_tensor * act1_alpha = nullptr;
|
||||
ggml_tensor * act1_beta = nullptr;
|
||||
ggml_tensor * conv1_w = nullptr;
|
||||
ggml_tensor * conv1_b = nullptr;
|
||||
ggml_tensor * act2_alpha = nullptr;
|
||||
ggml_tensor * act2_beta = nullptr;
|
||||
ggml_tensor * conv2_w = nullptr;
|
||||
ggml_tensor * conv2_b = nullptr;
|
||||
};
|
||||
|
||||
// one DAC upsample block (SnakeBeta -> causal ConvTranspose1d -> 3 residual units)
|
||||
struct dac_block {
|
||||
ggml_tensor * snake_alpha = nullptr;
|
||||
ggml_tensor * snake_beta = nullptr;
|
||||
ggml_tensor * conv_w = nullptr; // causal ConvTranspose1d
|
||||
ggml_tensor * conv_b = nullptr;
|
||||
std::vector<dac_res> res;
|
||||
};
|
||||
|
||||
// quantizer: RVQ codebook decode
|
||||
ggml_tensor * quant_first_in_w = nullptr; // semantic RVQ, in_proj (1x1 conv, loaded as 2D)
|
||||
ggml_tensor * quant_first_out_w = nullptr;
|
||||
ggml_tensor * quant_first_cb_w = nullptr; // codebook (1 layer)
|
||||
ggml_tensor * quant_rest_in_w = nullptr; // acoustic RVQ
|
||||
ggml_tensor * quant_rest_out_w = nullptr;
|
||||
ggml_tensor * quant_rest_cb_w = nullptr; // codebooks, merged 3D [15, vocab, dim]
|
||||
|
||||
ggml_tensor * pre_conv_w = nullptr;
|
||||
ggml_tensor * pre_conv_b = nullptr;
|
||||
|
||||
ggml_tensor * tfm_in_proj_w = nullptr;
|
||||
ggml_tensor * tfm_in_proj_b = nullptr;
|
||||
ggml_tensor * tfm_out_proj_w = nullptr;
|
||||
ggml_tensor * tfm_out_proj_b = nullptr;
|
||||
ggml_tensor * tfm_output_norm_w = nullptr;
|
||||
std::vector<clip_layer> tfm_layers; // reuses the generic block fields (ln_1/attn/ln_2/ffn/ls_1/ls_2)
|
||||
|
||||
std::vector<upsample_block> upsample;
|
||||
|
||||
ggml_tensor * dac_entry_w = nullptr;
|
||||
ggml_tensor * dac_entry_b = nullptr;
|
||||
std::vector<dac_block> dac;
|
||||
ggml_tensor * dac_post_snake_alpha = nullptr;
|
||||
ggml_tensor * dac_post_snake_beta = nullptr;
|
||||
ggml_tensor * dac_post_conv_w = nullptr;
|
||||
ggml_tensor * dac_post_conv_b = nullptr;
|
||||
};
|
||||
|
||||
struct clip_model {
|
||||
clip_modality modality = CLIP_MODALITY_VISION;
|
||||
projector_type proj_type = PROJECTOR_TYPE_MLP;
|
||||
@@ -577,6 +665,24 @@ struct clip_model {
|
||||
ggml_tensor * conv2d_3_w = nullptr;
|
||||
ggml_tensor * conv2d_3_b = nullptr;
|
||||
|
||||
// qwen3tts speaker encoder (ECAPA-TDNN)
|
||||
// reused tensors: stem conv is conv1d_1_w/b, feature aggregation is conv_out_w/b, output proj is mm_fc_w/b
|
||||
ggml_tensor * spk_asp_attn_w = nullptr;
|
||||
ggml_tensor * spk_asp_attn_b = nullptr;
|
||||
ggml_tensor * spk_asp_tdnn_w = nullptr;
|
||||
ggml_tensor * spk_asp_tdnn_b = nullptr;
|
||||
|
||||
// qwen3tts code_predictor
|
||||
ggml_tensor * gen_code_proj_in_w = nullptr; // small_to_mtp_projection
|
||||
ggml_tensor * gen_code_proj_in_b = nullptr;
|
||||
ggml_tensor * gen_code_embd_w = nullptr; // per-codebook embedding, merged 3D
|
||||
ggml_tensor * gen_code_head_w = nullptr; // per-codebook output head, merged 3D
|
||||
ggml_tensor * gen_code_out_embd_w = nullptr; // codebook-0 embedding, fed back into the talker
|
||||
ggml_tensor * gen_code_norm_w = nullptr; // final norm
|
||||
|
||||
// qwen3tts code2wav: RVQ codes -> raw PCM
|
||||
clip_code2wav c2w;
|
||||
|
||||
// cogvlm
|
||||
ggml_tensor * mm_post_fc_norm_w = nullptr;
|
||||
ggml_tensor * mm_post_fc_norm_b = nullptr;
|
||||
|
||||
+416
-39
@@ -17,6 +17,7 @@
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <map>
|
||||
#include <random>
|
||||
#include <stdexcept>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
@@ -269,6 +270,29 @@ clip_graph::clip_graph(clip_ctx * ctx, const clip_image_f32 & img) :
|
||||
gf = ggml_new_graph_custom(ctx0, ctx->max_nodes, false);
|
||||
}
|
||||
|
||||
clip_graph::clip_graph(const clip_graph & parent) :
|
||||
model(parent.model),
|
||||
hparams(parent.hparams),
|
||||
proj_type(parent.proj_type),
|
||||
img(parent.img),
|
||||
patch_size(parent.patch_size),
|
||||
n_patches_x(parent.n_patches_x),
|
||||
n_patches_y(parent.n_patches_y),
|
||||
n_patches(parent.n_patches),
|
||||
n_embd(parent.n_embd),
|
||||
n_head(parent.n_head),
|
||||
n_head_kv(parent.n_head_kv),
|
||||
d_head(parent.d_head),
|
||||
n_layer(parent.n_layer),
|
||||
n_mmproj_embd(parent.n_mmproj_embd),
|
||||
eps(parent.eps),
|
||||
kq_scale(parent.kq_scale),
|
||||
flash_attn_type(parent.flash_attn_type) {
|
||||
// reuse from parent
|
||||
ctx0 = parent.ctx0;
|
||||
gf = parent.gf;
|
||||
}
|
||||
|
||||
ggml_tensor * clip_graph::build_mm(ggml_tensor * w, ggml_tensor * x) const {
|
||||
return ggml_mul_mat(ctx0, w, x);
|
||||
}
|
||||
@@ -873,7 +897,8 @@ ggml_tensor * clip_graph::build_patch_merge_permute(ggml_tensor * cur, int scale
|
||||
return cur;
|
||||
}
|
||||
|
||||
static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const clip_image_f32_batch & imgs) {
|
||||
static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const clip_image_f32_batch & imgs,
|
||||
const clip_encode_params * params = nullptr) {
|
||||
const clip_image_f32 & img = imgs.entries[0];
|
||||
std::unique_ptr<clip_graph> builder;
|
||||
|
||||
@@ -1025,6 +1050,17 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
|
||||
{
|
||||
builder = std::make_unique<clip_graph_mimo_audio>(ctx, img);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_SPKENC:
|
||||
{
|
||||
builder = std::make_unique<clip_graph_qwen3tts_spkenc>(ctx, img);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
{
|
||||
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
|
||||
const int top_k = params ? params->top_k : 50;
|
||||
const float top_p = params ? params->top_p : 1.0f;
|
||||
builder = std::make_unique<clip_graph_qwen3tts_gen>(ctx, img, gen_process, top_k, top_p);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_YOUTUVL:
|
||||
{
|
||||
builder = std::make_unique<clip_graph_youtuvl>(ctx, img);
|
||||
@@ -1065,8 +1101,9 @@ struct clip_model_loader {
|
||||
|
||||
size_t model_size = 0; // in bytes
|
||||
|
||||
bool has_vision = false;
|
||||
bool has_audio = false;
|
||||
bool has_vision = false;
|
||||
bool has_audio = false;
|
||||
bool has_gen_audio = false;
|
||||
|
||||
mtmd_progress_callback progress_callback = nullptr;
|
||||
void * progress_callback_user_data = nullptr;
|
||||
@@ -1112,8 +1149,9 @@ struct clip_model_loader {
|
||||
|
||||
// modalities
|
||||
{
|
||||
get_bool(KEY_HAS_VISION_ENC, has_vision, false);
|
||||
get_bool(KEY_HAS_AUDIO_ENC, has_audio, false);
|
||||
get_bool(KEY_HAS_VISION_ENC, has_vision, false);
|
||||
get_bool(KEY_HAS_AUDIO_ENC, has_audio, false);
|
||||
get_bool(KEY_HAS_GEN_AUDIO_ENC, has_gen_audio, false);
|
||||
|
||||
if (has_vision) {
|
||||
LOG_INF("%s: has vision encoder\n", __func__);
|
||||
@@ -1121,6 +1159,9 @@ struct clip_model_loader {
|
||||
if (has_audio) {
|
||||
LOG_INF("%s: has audio encoder\n", __func__);
|
||||
}
|
||||
if (has_gen_audio) {
|
||||
LOG_INF("%s: has audio generation (gen) encoder\n", __func__);
|
||||
}
|
||||
}
|
||||
|
||||
// tensors
|
||||
@@ -1147,6 +1188,8 @@ struct clip_model_loader {
|
||||
GGML_ASSERT(has_vision);
|
||||
} else if (modality == CLIP_MODALITY_AUDIO) {
|
||||
GGML_ASSERT(has_audio);
|
||||
} else if (modality == CLIP_MODALITY_GEN_AUDIO) {
|
||||
GGML_ASSERT(has_gen_audio);
|
||||
}
|
||||
model.modality = modality;
|
||||
|
||||
@@ -1163,6 +1206,8 @@ struct clip_model_loader {
|
||||
get_string(KEY_VISION_PROJ_TYPE, proj_type, false);
|
||||
} else if (modality == CLIP_MODALITY_AUDIO) {
|
||||
get_string(KEY_AUDIO_PROJ_TYPE, proj_type, false);
|
||||
} else if (modality == CLIP_MODALITY_GEN_AUDIO) {
|
||||
get_string(KEY_GEN_AUDIO_PROJ_TYPE, proj_type, false);
|
||||
} else {
|
||||
GGML_ABORT("unknown modality");
|
||||
}
|
||||
@@ -1182,12 +1227,13 @@ struct clip_model_loader {
|
||||
}
|
||||
}
|
||||
|
||||
const bool is_vision = model.modality == CLIP_MODALITY_VISION;
|
||||
const bool is_audio = model.modality == CLIP_MODALITY_AUDIO;
|
||||
const bool is_vision = model.modality == CLIP_MODALITY_VISION;
|
||||
const bool is_audio = model.modality == CLIP_MODALITY_AUDIO;
|
||||
const bool is_gen_audio = model.modality == CLIP_MODALITY_GEN_AUDIO;
|
||||
|
||||
// other hparams
|
||||
{
|
||||
const char * prefix = is_vision ? "vision" : "audio";
|
||||
const char * prefix = is_vision ? "vision" : (is_audio ? "audio" : "gen.audio");
|
||||
get_u32(string_format(KEY_N_EMBD, prefix), hparams.n_embd);
|
||||
get_u32(string_format(KEY_N_HEAD, prefix), hparams.n_head);
|
||||
get_u32(string_format(KEY_N_EMBD_HEAD, prefix), hparams.n_embd_head, false);
|
||||
@@ -1198,6 +1244,7 @@ struct clip_model_loader {
|
||||
|
||||
// n_head_kv is optional (for GQA), default to n_head
|
||||
hparams.n_head_kv = hparams.n_head;
|
||||
get_u32(string_format(KEY_N_HEAD_KV, prefix), hparams.n_head_kv, false);
|
||||
|
||||
if (is_vision) {
|
||||
get_u32(KEY_IMAGE_SIZE, hparams.image_size);
|
||||
@@ -1226,6 +1273,11 @@ struct clip_model_loader {
|
||||
hparams.image_size = 0;
|
||||
hparams.patch_size = 1;
|
||||
|
||||
} else if (is_gen_audio) {
|
||||
// these are unused, but still need to be set to avoid issues
|
||||
hparams.image_size = 0;
|
||||
hparams.patch_size = 1;
|
||||
|
||||
} else {
|
||||
GGML_ASSERT(false && "unknown modality");
|
||||
}
|
||||
@@ -1647,6 +1699,33 @@ struct clip_model_loader {
|
||||
"%s: mimo_audio: %s must be > 0\n", __func__, KEY_A_LOCAL_GROUP_SIZE));
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_SPKENC:
|
||||
{
|
||||
// ECAPA-TDNN speaker encoder, mel front-end uses the Slaney default (fmin=0, fmax=sr/2)
|
||||
hparams.audio_sample_rate = 24000;
|
||||
hparams.audio_n_fft = 1024;
|
||||
hparams.audio_window_len = 1024;
|
||||
hparams.audio_hop_len = 256;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
{
|
||||
// TODO: hardcoded for now, read from code_predictor_config instead
|
||||
hparams.rope_theta = 1000000.0f;
|
||||
|
||||
// code2wav params
|
||||
hparams.wav_tfm_n_layer = 8;
|
||||
hparams.wav_tfm_n_embd = 512;
|
||||
hparams.wav_tfm_n_ff = 1024;
|
||||
hparams.wav_tfm_n_head = 16;
|
||||
hparams.wav_tfm_n_head_kv = 16;
|
||||
hparams.wav_tfm_eps = 1e-5f;
|
||||
hparams.wav_tfm_rope_theta = 10000.0f;
|
||||
hparams.wav_upsample_n_block = 2;
|
||||
hparams.wav_dac_n_block = 4;
|
||||
hparams.wav_dac_n_res = 3;
|
||||
// matches the reference decoder's sliding_window (speech_tokenizer/config.json)
|
||||
hparams.wav_tfm_swa = 72;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PADDLEOCR:
|
||||
{
|
||||
hparams.n_merge = 2;
|
||||
@@ -1871,7 +1950,9 @@ struct clip_model_loader {
|
||||
}
|
||||
|
||||
// TODO @ngxson : support both audio and video in the future
|
||||
const char * prefix = model.modality == CLIP_MODALITY_AUDIO ? "a" : "v";
|
||||
const char * prefix = model.modality == CLIP_MODALITY_AUDIO ? "a"
|
||||
: model.modality == CLIP_MODALITY_GEN_AUDIO ? "a.gen.code"
|
||||
: "v";
|
||||
|
||||
// get offsets
|
||||
for (int64_t i = 0; i < gguf_get_n_tensors(ctx_gguf.get()); ++i) {
|
||||
@@ -1973,7 +2054,8 @@ struct clip_model_loader {
|
||||
model.position_embeddings = get_tensor(string_format(TN_POS_EMBD, prefix), false);
|
||||
|
||||
const bool has_standard_layers = (
|
||||
model.proj_type != PROJECTOR_TYPE_GEMMA3NV);
|
||||
model.proj_type != PROJECTOR_TYPE_GEMMA3NV &&
|
||||
model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC);
|
||||
|
||||
// layers
|
||||
const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0;
|
||||
@@ -2599,6 +2681,144 @@ struct clip_model_loader {
|
||||
model.mm_1_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "weight"));
|
||||
model.mm_2_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "weight"));
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_SPKENC:
|
||||
{
|
||||
// stem TDNN (block 0)
|
||||
model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 0, "weight"));
|
||||
model.conv1d_1_b = get_tensor(string_format(TN_CONV1D, 0, "bias"));
|
||||
|
||||
// SE-Res2Net blocks (GGUF bid 1..3, one per hparams.n_layer)
|
||||
model.layers.resize(hparams.n_layer);
|
||||
for (int il = 0; il < hparams.n_layer; il++) {
|
||||
auto & layer = model.layers[il];
|
||||
int bid = il + 1;
|
||||
layer.conv_pw1_w = get_tensor(string_format(TN_CONV_PW1, prefix, bid, "weight"));
|
||||
layer.conv_pw1_b = get_tensor(string_format(TN_CONV_PW1, prefix, bid, "bias"));
|
||||
layer.conv_pw2_w = get_tensor(string_format(TN_CONV_PW2, prefix, bid, "weight"));
|
||||
layer.conv_pw2_b = get_tensor(string_format(TN_CONV_PW2, prefix, bid, "bias"));
|
||||
layer.se_conv1_w = get_tensor(string_format(TN_A_SE_CONV1, bid, "weight"));
|
||||
layer.se_conv1_b = get_tensor(string_format(TN_A_SE_CONV1, bid, "bias"));
|
||||
layer.se_conv2_w = get_tensor(string_format(TN_A_SE_CONV2, bid, "weight"));
|
||||
layer.se_conv2_b = get_tensor(string_format(TN_A_SE_CONV2, bid, "bias"));
|
||||
layer.res2_conv_w.resize(7);
|
||||
layer.res2_conv_b.resize(7);
|
||||
for (int xid = 0; xid < 7; xid++) {
|
||||
layer.res2_conv_w[xid] = get_tensor(string_format(TN_A_CONV_RES2, bid, xid, "weight"));
|
||||
layer.res2_conv_b[xid] = get_tensor(string_format(TN_A_CONV_RES2, bid, xid, "bias"));
|
||||
}
|
||||
}
|
||||
|
||||
// multi-layer feature aggregation
|
||||
model.conv_out_w = get_tensor(string_format(TN_CONV_OUT, "weight"));
|
||||
model.conv_out_b = get_tensor(string_format(TN_CONV_OUT, "bias"));
|
||||
|
||||
// attentive statistics pooling
|
||||
model.spk_asp_attn_w = get_tensor(string_format(TN_A_ASP_ATTN, "weight"));
|
||||
model.spk_asp_attn_b = get_tensor(string_format(TN_A_ASP_ATTN, "bias"));
|
||||
model.spk_asp_tdnn_w = get_tensor(string_format(TN_A_ASP_TDNN, "weight"));
|
||||
model.spk_asp_tdnn_b = get_tensor(string_format(TN_A_ASP_TDNN, "bias"));
|
||||
|
||||
// final speaker embedding projection
|
||||
model.mm_fc_w = get_tensor(string_format(TN_MM_AUDIO_FC, "weight"));
|
||||
model.mm_fc_b = get_tensor(string_format(TN_MM_AUDIO_FC, "bias"));
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
{
|
||||
// code_predictor
|
||||
model.gen_code_proj_in_w = get_tensor(string_format(TN_A_GEN_CODE_PROJ_IN, "weight"));
|
||||
model.gen_code_proj_in_b = get_tensor(string_format(TN_A_GEN_CODE_PROJ_IN, "bias"));
|
||||
model.gen_code_embd_w = get_tensor(string_format(TN_A_GEN_CODE_EMBD, "weight"));
|
||||
model.gen_code_head_w = get_tensor(string_format(TN_A_GEN_CODE_HEAD, "weight"));
|
||||
model.gen_code_out_embd_w = get_tensor(string_format(TN_A_GEN_CODE_OUT_EMBD, "weight"));
|
||||
model.gen_code_norm_w = get_tensor(string_format(TN_A_GEN_CODE_NORM, "weight"));
|
||||
|
||||
// code2wav: RVQ codes -> raw PCM, lives in the same ctx as code_predictor
|
||||
{
|
||||
auto & c2w = model.c2w;
|
||||
|
||||
c2w.quant_first_in_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_IN, "weight"));
|
||||
c2w.quant_first_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_OUT, "weight"));
|
||||
c2w.quant_first_cb_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_CB, "weight"));
|
||||
c2w.quant_rest_in_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_REST_IN, "weight"));
|
||||
c2w.quant_rest_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_REST_OUT, "weight"));
|
||||
c2w.quant_rest_cb_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_REST_CB, "weight"));
|
||||
|
||||
c2w.pre_conv_w = get_tensor(string_format(TN_A_GEN_WAV_PRE_CONV, "weight"));
|
||||
c2w.pre_conv_b = get_tensor(string_format(TN_A_GEN_WAV_PRE_CONV, "bias"));
|
||||
|
||||
c2w.tfm_in_proj_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_IN_PROJ, "weight"));
|
||||
c2w.tfm_in_proj_b = get_tensor(string_format(TN_A_GEN_WAV_TFM_IN_PROJ, "bias"));
|
||||
c2w.tfm_out_proj_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_PROJ, "weight"));
|
||||
c2w.tfm_out_proj_b = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_PROJ, "bias"));
|
||||
c2w.tfm_output_norm_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_NORM, "weight"));
|
||||
|
||||
// loaded manually, the generic model.layers loop is taken by code_predictor
|
||||
c2w.tfm_layers.resize(hparams.wav_tfm_n_layer);
|
||||
for (int il = 0; il < hparams.wav_tfm_n_layer; il++) {
|
||||
auto & layer = c2w.tfm_layers[il];
|
||||
const char * p = "a.gen.wav.tfm";
|
||||
layer.q_w = get_tensor(string_format(TN_ATTN_Q, p, il, "weight"));
|
||||
layer.k_w = get_tensor(string_format(TN_ATTN_K, p, il, "weight"));
|
||||
layer.v_w = get_tensor(string_format(TN_ATTN_V, p, il, "weight"));
|
||||
layer.o_w = get_tensor(string_format(TN_ATTN_OUTPUT, p, il, "weight"));
|
||||
layer.ln_1_w = get_tensor(string_format(TN_LN_1, p, il, "weight"));
|
||||
layer.ln_2_w = get_tensor(string_format(TN_LN_2, p, il, "weight"));
|
||||
layer.ls_1_w = get_tensor(string_format(TN_LS_1, p, il, "weight"));
|
||||
layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight"));
|
||||
layer.ff_gate_w = get_tensor(string_format(TN_FFN_GATE, p, il, "weight"));
|
||||
layer.ff_up_w = get_tensor(string_format(TN_FFN_UP, p, il, "weight"));
|
||||
layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, p, il, "weight"));
|
||||
}
|
||||
|
||||
// upsample: 2x (causal ConvTranspose1d + ConvNeXt block)
|
||||
c2w.upsample.resize(hparams.wav_upsample_n_block);
|
||||
for (int il = 0; il < hparams.wav_upsample_n_block; il++) {
|
||||
auto & up = c2w.upsample[il];
|
||||
up.conv_w = get_tensor(string_format(TN_A_GEN_WAV_UP_CONV, il, "weight"));
|
||||
up.conv_b = get_tensor(string_format(TN_A_GEN_WAV_UP_CONV, il, "bias"));
|
||||
up.dwconv_w = get_tensor(string_format(TN_A_GEN_WAV_UP_DWCONV, il, "weight"));
|
||||
up.dwconv_b = get_tensor(string_format(TN_A_GEN_WAV_UP_DWCONV, il, "bias"));
|
||||
up.norm_w = get_tensor(string_format(TN_A_GEN_WAV_UP_NORM, il, "weight"));
|
||||
up.norm_b = get_tensor(string_format(TN_A_GEN_WAV_UP_NORM, il, "bias"));
|
||||
up.pw1_w = get_tensor(string_format(TN_A_GEN_WAV_UP_PW1, il, "weight"));
|
||||
up.pw1_b = get_tensor(string_format(TN_A_GEN_WAV_UP_PW1, il, "bias"));
|
||||
up.pw2_w = get_tensor(string_format(TN_A_GEN_WAV_UP_PW2, il, "weight"));
|
||||
up.pw2_b = get_tensor(string_format(TN_A_GEN_WAV_UP_PW2, il, "bias"));
|
||||
up.gamma = get_tensor(string_format(TN_A_GEN_WAV_UP_GAMMA, il));
|
||||
}
|
||||
|
||||
// DAC decoder: conv_pre + n upsample blocks (each with n_res residual units) + conv_post
|
||||
c2w.dac_entry_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_ENTRY, "weight"));
|
||||
c2w.dac_entry_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_ENTRY, "bias"));
|
||||
|
||||
c2w.dac.resize(hparams.wav_dac_n_block);
|
||||
for (int il = 0; il < hparams.wav_dac_n_block; il++) {
|
||||
auto & blk = c2w.dac[il];
|
||||
blk.snake_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_SNAKE, il, "alpha"));
|
||||
blk.snake_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_SNAKE, il, "beta"));
|
||||
blk.conv_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_CONV, il, "weight"));
|
||||
blk.conv_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_CONV, il, "bias"));
|
||||
|
||||
blk.res.resize(hparams.wav_dac_n_res);
|
||||
for (int ir = 0; ir < hparams.wav_dac_n_res; ir++) {
|
||||
auto & res = blk.res[ir];
|
||||
res.act1_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT1, il, ir, "alpha"));
|
||||
res.act1_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT1, il, ir, "beta"));
|
||||
res.conv1_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV1, il, ir, "weight"));
|
||||
res.conv1_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV1, il, ir, "bias"));
|
||||
res.act2_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT2, il, ir, "alpha"));
|
||||
res.act2_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT2, il, ir, "beta"));
|
||||
res.conv2_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV2, il, ir, "weight"));
|
||||
res.conv2_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV2, il, ir, "bias"));
|
||||
}
|
||||
}
|
||||
|
||||
c2w.dac_post_snake_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_SNAKE, "alpha"));
|
||||
c2w.dac_post_snake_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_SNAKE, "beta"));
|
||||
c2w.dac_post_conv_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_CONV, "weight"));
|
||||
c2w.dac_post_conv_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_CONV, "bias"));
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_VOXTRAL:
|
||||
{
|
||||
model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 1, "weight"));
|
||||
@@ -3427,6 +3647,7 @@ struct clip_model_loader {
|
||||
struct clip_init_result clip_init(const char * fname, struct clip_context_params ctx_params) {
|
||||
clip_ctx * ctx_vision = nullptr;
|
||||
clip_ctx * ctx_audio = nullptr;
|
||||
clip_ctx * ctx_gen_audio = nullptr;
|
||||
|
||||
try {
|
||||
clip_model_loader loader(fname,
|
||||
@@ -3459,16 +3680,25 @@ struct clip_init_result clip_init(const char * fname, struct clip_context_params
|
||||
}
|
||||
}
|
||||
|
||||
if (loader.has_gen_audio) {
|
||||
ctx_gen_audio = new clip_ctx(ctx_params);
|
||||
loader.load_hparams(ctx_gen_audio->model, CLIP_MODALITY_GEN_AUDIO);
|
||||
loader.load_tensors(*ctx_gen_audio);
|
||||
// TODO: fix warmup
|
||||
ctx_gen_audio->buf_compute_meta.resize(ctx_gen_audio->max_nodes * ggml_tensor_overhead() + ggml_graph_overhead());
|
||||
}
|
||||
|
||||
} catch (const std::exception & e) {
|
||||
LOG_ERR("%s: failed to load model '%s': %s\n", __func__, fname, e.what());
|
||||
|
||||
delete ctx_vision;
|
||||
delete ctx_audio;
|
||||
delete ctx_gen_audio;
|
||||
|
||||
return {nullptr, nullptr};
|
||||
return {nullptr, nullptr, nullptr};
|
||||
}
|
||||
|
||||
return {ctx_vision, ctx_audio};
|
||||
return {ctx_vision, ctx_audio, ctx_gen_audio};
|
||||
}
|
||||
|
||||
struct clip_cap clip_get_cap(const char * fname) {
|
||||
@@ -3784,6 +4014,16 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
|
||||
const int ds = ctx->model.hparams.audio_proj_downsample_rate;
|
||||
n_patches = ((img->nx() + ws - 1) / ws) * (ws / ds);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_SPKENC:
|
||||
{
|
||||
// pooling gives one speaker embedding, whatever the clip length is
|
||||
n_patches = 1;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
{
|
||||
// one hidden-state vector fed back to the talker per call
|
||||
n_patches = 1;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GRANITE4_VISION:
|
||||
{
|
||||
// Per-tile output token count: each projector block outputs
|
||||
@@ -3817,7 +4057,16 @@ bool clip_image_encode(struct clip_ctx * ctx, int n_threads, const clip_image_f3
|
||||
}
|
||||
|
||||
bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32_batch * imgs_c_ptr, std::vector<float> & out_batch_embd) {
|
||||
const clip_image_f32_batch & imgs = *imgs_c_ptr;
|
||||
clip_encode_params params;
|
||||
params.imgs = imgs_c_ptr;
|
||||
params.n_threads = n_threads;
|
||||
params.out_embd = &out_batch_embd;
|
||||
|
||||
return clip_encode(ctx, ¶ms);
|
||||
}
|
||||
|
||||
bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
const clip_image_f32_batch & imgs = *params->imgs;
|
||||
int n_batch_cur = imgs.entries.size();
|
||||
|
||||
// [QWEN_VIDEO] for video models, the batch dimension is used as temporal dimension for merged frames
|
||||
@@ -3828,12 +4077,12 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
|
||||
|
||||
// if buffers are not allocated, we need to do a warmup run to allocate them
|
||||
if (!ctx->is_allocated) {
|
||||
clip_model_loader::warmup(*ctx, *imgs_c_ptr);
|
||||
clip_model_loader::warmup(*ctx, *params->imgs);
|
||||
}
|
||||
|
||||
// build the inference graph
|
||||
ggml_backend_sched_reset(ctx->sched.get());
|
||||
ggml_cgraph * gf = clip_get_graph_builder(ctx, imgs)->build();
|
||||
ggml_cgraph * gf = clip_get_graph_builder(ctx, imgs, params)->build();
|
||||
ggml_backend_sched_alloc_graph(ctx->sched.get(), gf);
|
||||
|
||||
// set inputs
|
||||
@@ -3918,8 +4167,8 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
|
||||
}
|
||||
set_input_f32("inp_raw", inp_raw);
|
||||
|
||||
} else {
|
||||
// audio input
|
||||
} else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_GEN_WAV)) {
|
||||
// audio input, code2wav is not here: its only input is "inp_codes", set in the switch below
|
||||
GGML_ASSERT(imgs.entries.size() == 1);
|
||||
|
||||
const auto & mel_inp = imgs.entries[0];
|
||||
@@ -4475,9 +4724,77 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
|
||||
case PROJECTOR_TYPE_COGVLM:
|
||||
case PROJECTOR_TYPE_YASA2:
|
||||
case PROJECTOR_TYPE_GEMMA4UA:
|
||||
case PROJECTOR_TYPE_QWEN3TTS_SPKENC:
|
||||
{
|
||||
// do nothing
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
{
|
||||
if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) {
|
||||
GGML_ASSERT(params->codes != nullptr);
|
||||
|
||||
// frame-major input to group-major, rear-padded with code 0 up to one window
|
||||
const int64_t n_codes = model.gen_code_head_w->ne[2] + 1;
|
||||
const int64_t n_frames_w = hparams.wav_tfm_swa;
|
||||
const int64_t n_frames = (int64_t) params->codes->size() / n_codes;
|
||||
GGML_ASSERT(n_frames > 0 && n_frames <= n_frames_w);
|
||||
|
||||
// codes are used as ggml_get_rows indices, so check them against the codebook vocab
|
||||
const int64_t vocab_first = model.c2w.quant_first_cb_w->ne[1];
|
||||
const int64_t vocab_rest = model.c2w.quant_rest_cb_w->ne[1];
|
||||
for (int64_t f = 0; f < n_frames; f++) {
|
||||
for (int64_t g = 0; g < n_codes; g++) {
|
||||
const int32_t c = (*params->codes)[f * n_codes + g];
|
||||
const int64_t vocab = (g == 0) ? vocab_first : vocab_rest;
|
||||
if (c < 0 || (int64_t) c >= vocab) {
|
||||
LOG_ERR("%s: code out of range (frame %lld, group %lld, code %d, vocab %lld)\n",
|
||||
__func__, (long long) f, (long long) g, c, (long long) vocab);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int32_t> codes(n_frames_w * n_codes, 0);
|
||||
for (int64_t f = 0; f < n_frames; f++) {
|
||||
for (int64_t g = 0; g < n_codes; g++) {
|
||||
codes[g * n_frames_w + f] = (*params->codes)[f * n_codes + g];
|
||||
}
|
||||
}
|
||||
set_input_i32("inp_codes", codes);
|
||||
|
||||
// upload the state from the previous call, or zero-fill on a cold start
|
||||
size_t offset = 0;
|
||||
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
|
||||
ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str());
|
||||
const size_t nb = ggml_nbytes(t);
|
||||
if (params->state_in && params->state_in->size() >= offset + nb) {
|
||||
ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb);
|
||||
} else {
|
||||
std::vector<uint8_t> zeros(nb, 0);
|
||||
ggml_backend_tensor_set(t, zeros.data(), 0, nb);
|
||||
}
|
||||
offset += nb;
|
||||
}
|
||||
} else {
|
||||
// code0 indexes gen_code_out_embd_w via ggml_get_rows; bound it
|
||||
const int64_t vocab0 = model.gen_code_out_embd_w->ne[1];
|
||||
if (params->code0 < 0 || (int64_t) params->code0 >= vocab0) {
|
||||
LOG_ERR("%s: code0 out of range (%d, vocab %lld)\n", __func__, params->code0, (long long) vocab0);
|
||||
return false;
|
||||
}
|
||||
std::vector<int32_t> code0 = { params->code0 };
|
||||
set_input_i32("inp_code0", code0);
|
||||
|
||||
// one uniform(0,1) draw per codebook, used by do_sampling()
|
||||
static std::mt19937 rng{ std::random_device{}() };
|
||||
std::uniform_real_distribution<float> dist(0.0f, 1.0f);
|
||||
const int64_t n_acoustic = model.gen_code_head_w->ne[2];
|
||||
for (int64_t g = 0; g < n_acoustic; g++) {
|
||||
std::vector<float> r = { dist(rng) };
|
||||
set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r);
|
||||
}
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_HUNYUANVL:
|
||||
{
|
||||
// Compute the HunyuanVL 2D position embedding on CPU (with the
|
||||
@@ -4883,7 +5200,7 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
|
||||
if (reg) {
|
||||
auto ggml_backend_set_n_threads_fn = (ggml_backend_set_n_threads_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads");
|
||||
if (ggml_backend_set_n_threads_fn) {
|
||||
ggml_backend_set_n_threads_fn(ctx->backend_cpu, n_threads);
|
||||
ggml_backend_set_n_threads_fn(ctx->backend_cpu, params->n_threads);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4893,34 +5210,90 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
|
||||
return false;
|
||||
}
|
||||
|
||||
// the last node is the embedding tensor
|
||||
ggml_tensor * embeddings = ggml_graph_node(gf, -1);
|
||||
// the last node is the embedding tensor, code2wav has no out_embd
|
||||
ggml_tensor * embeddings = params->out_embd ? ggml_graph_node(gf, -1) : nullptr;
|
||||
|
||||
// sanity check (assuming that all images in batch have the same number of tokens, so we only check the first one)
|
||||
const int n_tokens_out = embeddings->ne[1];
|
||||
const int expected_n_tokens_out = clip_n_output_tokens(ctx, &imgs.entries[0]);
|
||||
if (n_tokens_out != expected_n_tokens_out) {
|
||||
LOG_ERR("%s: expected output %d tokens, got %d\n", __func__, expected_n_tokens_out, n_tokens_out);
|
||||
GGML_ABORT("Invalid number of output tokens");
|
||||
}
|
||||
if (embeddings != nullptr) {
|
||||
// sanity check (assuming that all images in batch have the same number of tokens, so we only check the first one)
|
||||
const int n_tokens_out = embeddings->ne[1];
|
||||
const int expected_n_tokens_out = clip_n_output_tokens(ctx, &imgs.entries[0]);
|
||||
if (n_tokens_out != expected_n_tokens_out) {
|
||||
LOG_ERR("%s: expected output %d tokens, got %d\n", __func__, expected_n_tokens_out, n_tokens_out);
|
||||
GGML_ABORT("Invalid number of output tokens");
|
||||
}
|
||||
|
||||
LOG_DBG("%s: output embedding shape [%d, %d, %d]\n", __func__,
|
||||
(int)embeddings->ne[0], (int)embeddings->ne[1], (int)embeddings->ne[2]);
|
||||
LOG_DBG("%s: output embedding shape [%d, %d, %d]\n", __func__,
|
||||
(int)embeddings->ne[0], (int)embeddings->ne[1], (int)embeddings->ne[2]);
|
||||
|
||||
// copy output to user buffer if provided
|
||||
// if output is empty, skip the copy
|
||||
if (!out_batch_embd.empty()) {
|
||||
if (out_batch_embd.size() != (size_t)ggml_nelements(embeddings)) {
|
||||
LOG_ERR("%s: output buffer has %zu elements but expected %zu\n", __func__, out_batch_embd.size(), (size_t)ggml_nelements(embeddings));
|
||||
GGML_ABORT("Output buffer size mismatch");
|
||||
// copy output to user buffer if provided
|
||||
// if output is empty, skip the copy
|
||||
auto & out_batch_embd = *params->out_embd;
|
||||
if (!out_batch_embd.empty()) {
|
||||
if (out_batch_embd.size() != (size_t)ggml_nelements(embeddings)) {
|
||||
LOG_ERR("%s: output buffer has %zu elements but expected %zu\n", __func__, out_batch_embd.size(), (size_t)ggml_nelements(embeddings));
|
||||
GGML_ABORT("Output buffer size mismatch");
|
||||
}
|
||||
ggml_backend_tensor_get(embeddings, out_batch_embd.data(), 0, ggml_nbytes(embeddings));
|
||||
} else {
|
||||
LOG_WRN("%s: output buffer is empty, skipping copy\n", __func__);
|
||||
}
|
||||
ggml_backend_tensor_get(embeddings, out_batch_embd.data(), 0, ggml_nbytes(embeddings));
|
||||
} else {
|
||||
LOG_WRN("%s: output buffer is empty, skipping copy\n", __func__);
|
||||
}
|
||||
|
||||
//
|
||||
// for audio gen models
|
||||
//
|
||||
|
||||
if (params->out_codes != nullptr) {
|
||||
ggml_tensor * codes = ggml_graph_get_tensor(gf, "out_codes");
|
||||
if (codes == nullptr) {
|
||||
GGML_ABORT("out_codes requested but graph has no \"out_codes\" tensor");
|
||||
}
|
||||
auto & out_codes = *params->out_codes;
|
||||
out_codes.resize(ggml_nelements(codes));
|
||||
ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes));
|
||||
}
|
||||
if (params->out_audio != nullptr) {
|
||||
ggml_tensor * audio = ggml_graph_get_tensor(gf, "out_audio");
|
||||
if (audio == nullptr) {
|
||||
GGML_ABORT("out_audio requested but graph has no \"out_audio\" tensor");
|
||||
}
|
||||
auto & out_audio = *params->out_audio;
|
||||
out_audio.resize(ggml_nelements(audio));
|
||||
ggml_backend_tensor_get(audio, out_audio.data(), 0, ggml_nbytes(audio));
|
||||
|
||||
// drop the tail audio that comes from the code-0 rear padding
|
||||
const int64_t n_codes = model.gen_code_head_w->ne[2] + 1;
|
||||
const int64_t n_frames_w = hparams.wav_tfm_swa;
|
||||
const int64_t n_frames = (int64_t) params->codes->size() / n_codes;
|
||||
if (n_frames < n_frames_w) {
|
||||
const size_t hop = out_audio.size() / n_frames_w;
|
||||
out_audio.resize((size_t) n_frames * hop);
|
||||
}
|
||||
}
|
||||
if (params->state_out != nullptr) {
|
||||
auto & state_out = *params->state_out;
|
||||
size_t total = 0;
|
||||
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
|
||||
total += (size_t) (slot.ne0 * slot.ne1) * sizeof(float);
|
||||
}
|
||||
state_out.resize(total);
|
||||
size_t offset = 0;
|
||||
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
|
||||
ggml_tensor * t = ggml_graph_get_tensor(gf, ("state_out_" + slot.name).c_str());
|
||||
if (t == nullptr) {
|
||||
GGML_ABORT("state_out requested but graph has no \"state_out_%s\" tensor", slot.name.c_str());
|
||||
}
|
||||
const size_t nb = ggml_nbytes(t);
|
||||
ggml_backend_tensor_get(t, state_out.data() + offset, 0, nb);
|
||||
offset += nb;
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Debug: dump final embeddings if MTMD_DEBUG_EMBEDDINGS is set
|
||||
if (ctx->debug_output_embeddings) {
|
||||
//
|
||||
|
||||
if (ctx->debug_output_embeddings && embeddings != nullptr) {
|
||||
const int64_t n_embd = embeddings->ne[0];
|
||||
const int64_t n_tokens = embeddings->ne[1];
|
||||
std::vector<float> emb_data(ggml_nelements(embeddings));
|
||||
@@ -5047,6 +5420,10 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
|
||||
return ctx->model.mm_ffn_down_w->ne[1];
|
||||
case PROJECTOR_TYPE_MIMO_AUDIO:
|
||||
return ctx->model.mm_2_w->ne[1];
|
||||
case PROJECTOR_TYPE_QWEN3TTS_SPKENC:
|
||||
return ctx->model.mm_fc_w->ne[2];
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
return ctx->model.gen_code_out_embd_w->ne[0];
|
||||
case PROJECTOR_TYPE_PARAKEET:
|
||||
return ctx->model.mm_1_w->ne[1];
|
||||
default:
|
||||
|
||||
@@ -37,6 +37,7 @@ struct clip_image_f32_batch;
|
||||
enum clip_modality {
|
||||
CLIP_MODALITY_VISION,
|
||||
CLIP_MODALITY_AUDIO,
|
||||
CLIP_MODALITY_GEN_AUDIO,
|
||||
};
|
||||
|
||||
enum clip_flash_attn_type {
|
||||
@@ -61,6 +62,7 @@ struct clip_context_params {
|
||||
struct clip_init_result {
|
||||
struct clip_ctx * ctx_v; // vision context
|
||||
struct clip_ctx * ctx_a; // audio context
|
||||
struct clip_ctx * ctx_gen_a; // audio generation context
|
||||
};
|
||||
|
||||
struct clip_init_result clip_init(const char * fname, struct clip_context_params ctx_params);
|
||||
@@ -84,6 +86,33 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx);
|
||||
bool clip_image_encode (struct clip_ctx * ctx, int n_threads, const clip_image_f32 * img, std::vector<float> & out_vec);
|
||||
bool clip_image_batch_encode(struct clip_ctx * ctx, int n_threads, const struct clip_image_f32_batch * imgs, std::vector<float> & out_batch_embd);
|
||||
|
||||
enum clip_gen_process_type {
|
||||
CLIP_GEN_PROCESS_GEN_UNKNOWN,
|
||||
CLIP_GEN_PROCESS_GEN_CODE, // h_state to codes
|
||||
CLIP_GEN_PROCESS_GEN_WAV, // codes to raw PCM audio
|
||||
};
|
||||
struct clip_encode_params {
|
||||
int n_threads = 1;
|
||||
const clip_image_f32_batch * imgs = nullptr;
|
||||
std::vector<float> * out_embd = nullptr;
|
||||
|
||||
// for audio gen, imgs has exactly one entry: hidden state from backbone (GEN_CODE) or unused (GEN_WAV)
|
||||
clip_gen_process_type gen_process = CLIP_GEN_PROCESS_GEN_UNKNOWN;
|
||||
|
||||
// GEN_CODE: out_embd receives the embd to feed back to the backbone
|
||||
int32_t code0 = 0; // semantic code sampled by the backbone
|
||||
int32_t top_k = 50;
|
||||
float top_p = 1.0f;
|
||||
std::vector<int32_t> * out_codes = nullptr; // this frame's 16 sampled codes
|
||||
|
||||
// GEN_WAV
|
||||
const std::vector<int32_t> * codes = nullptr; // this frame's 16 RVQ codes
|
||||
std::vector<float> * out_audio = nullptr; // decoded PCM samples, F32
|
||||
const std::vector<uint8_t> * state_in = nullptr; // state from previous call, null or wrong size means cold start
|
||||
std::vector<uint8_t> * state_out = nullptr; // state for the next call
|
||||
};
|
||||
bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params);
|
||||
|
||||
bool clip_is_llava(const struct clip_ctx * ctx);
|
||||
// note for contributor: this clip_is_(model) pattern is deprecated
|
||||
// do NOT add new functions like this
|
||||
|
||||
@@ -2,6 +2,11 @@
|
||||
|
||||
#include "../clip-graph.h"
|
||||
|
||||
#include <map>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
/*
|
||||
* IMPORTANT: The mtmd module does NOT accept pull requests that are fully or predominantly AI-generated.
|
||||
* We encourage human contributors to ensure the quality and reliability of the codebase.
|
||||
@@ -215,6 +220,111 @@ struct clip_graph_mimo_audio : clip_graph {
|
||||
ggml_cgraph * build() override;
|
||||
};
|
||||
|
||||
struct clip_graph_qwen3tts_spkenc : clip_graph {
|
||||
clip_graph_qwen3tts_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
ggml_tensor * conv1d_same(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int dilation) const;
|
||||
ggml_tensor * res2net(ggml_tensor * x, const clip_layer & layer, int dilation, int scale) const;
|
||||
ggml_tensor * se_block(ggml_tensor * x, const clip_layer & layer) const;
|
||||
ggml_tensor * se_res2net_block(ggml_tensor * x, const clip_layer & layer, int dilation, int scale) const;
|
||||
ggml_tensor * attentive_stats_pool(ggml_tensor * x) const;
|
||||
};
|
||||
|
||||
struct clip_graph_qwen3tts_gen : clip_graph {
|
||||
clip_graph_qwen3tts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int top_k, float top_p)
|
||||
: clip_graph(ctx, img), gen_process(gen_process), top_k(top_k), top_p(top_p) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
// which sub-graph build() constructs, fixed at graph-build time
|
||||
clip_gen_process_type gen_process;
|
||||
|
||||
// sampling params, fixed at graph-build time (GEN_CODE only)
|
||||
int top_k;
|
||||
float top_p;
|
||||
|
||||
//
|
||||
// code_gen: backbone hidden state + sampled code0 -> 16 RVQ codes
|
||||
// MTP-style code predictor, one token per codebook
|
||||
//
|
||||
struct code_gen : clip_graph {
|
||||
code_gen(const clip_graph & parent, int top_k, float top_p)
|
||||
: clip_graph(parent), top_k(top_k), top_p(top_p) {}
|
||||
ggml_cgraph * build() override { GGML_ABORT("call prefill()/step() instead"); }
|
||||
|
||||
int top_k;
|
||||
float top_p;
|
||||
|
||||
ggml_tensor * cache_set(ggml_tensor * cache, int row_idx, ggml_tensor * value) const;
|
||||
ggml_tensor * do_sampling(ggml_tensor * logits, ggml_tensor * inp_rand) const;
|
||||
|
||||
ggml_tensor * const_i32(ggml_tensor * anchor, float value) const;
|
||||
ggml_tensor * causal_mask_row(int64_t n_kv_pad, int pos) const;
|
||||
ggml_tensor * project_in(ggml_tensor * cur) const;
|
||||
|
||||
ggml_tensor * layer_forward(
|
||||
ggml_tensor * cur,
|
||||
const clip_layer & layer,
|
||||
ggml_tensor * inp_pos,
|
||||
ggml_tensor * kq_mask,
|
||||
ggml_tensor *& k_cache_layer,
|
||||
ggml_tensor *& v_cache_layer,
|
||||
int64_t n_kv_pad,
|
||||
int pos,
|
||||
int il) const;
|
||||
|
||||
void prefill(
|
||||
std::vector<ggml_tensor *> & k_cache,
|
||||
std::vector<ggml_tensor *> & v_cache,
|
||||
ggml_tensor *& out_code_cache,
|
||||
ggml_tensor * h_state,
|
||||
ggml_tensor * code0_embd,
|
||||
ggml_tensor * inp_rand) const;
|
||||
|
||||
ggml_tensor * step(
|
||||
std::vector<ggml_tensor *> & k_cache,
|
||||
std::vector<ggml_tensor *> & v_cache,
|
||||
ggml_tensor * out_code_cache,
|
||||
ggml_tensor * inp_rand,
|
||||
int step_idx) const;
|
||||
};
|
||||
|
||||
//
|
||||
// code2wav: RVQ codes -> raw PCM (quantizer + pre_conv + pre_transformer + upsample + DAC).
|
||||
//
|
||||
struct code2wav : clip_graph {
|
||||
code2wav(const clip_graph & parent) : clip_graph(parent) {}
|
||||
ggml_cgraph * build() override { GGML_ABORT("call decode() instead"); }
|
||||
|
||||
// state_in: previous call's persisted state, by slot name (see list_c2w_state_slots())
|
||||
std::map<std::string, ggml_tensor *> state_in;
|
||||
// state_out: this call's state to persist, added to the graph outputs by build()
|
||||
mutable std::vector<std::pair<std::string, ggml_tensor *>> state_out;
|
||||
|
||||
// stateful conv ops: read/update their state via state_in/state_out[state_name]
|
||||
ggml_tensor * causal_conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int dilation, const std::string & state_name) const;
|
||||
ggml_tensor * causal_conv1d_dw(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, const std::string & state_name) const;
|
||||
ggml_tensor * causal_conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, const std::string & state_name) const;
|
||||
ggml_tensor * snake(ggml_tensor * x, ggml_tensor * alpha, ggml_tensor * beta) const;
|
||||
|
||||
ggml_tensor * quant_decode(ggml_tensor * inp_codes) const;
|
||||
ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, int il) const;
|
||||
ggml_tensor * convnext_block(ggml_tensor * x, const clip_code2wav::upsample_block & blk, const std::string & state_prefix) const;
|
||||
ggml_tensor * dac_res_unit(ggml_tensor * x, const clip_code2wav::dac_res & res, int dilation, const std::string & state_name) const;
|
||||
|
||||
// inp_codes [1, n_codes] I32 -> this frame's audio samples [n_samples] F32, clamped to [-1, 1]
|
||||
ggml_tensor * decode(ggml_tensor * inp_codes) const;
|
||||
};
|
||||
};
|
||||
|
||||
// one persisted state buffer used by code2wav, see qwen3tts-gen.cpp
|
||||
struct c2w_state_slot {
|
||||
std::string name;
|
||||
int64_t ne0;
|
||||
int64_t ne1;
|
||||
};
|
||||
std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model);
|
||||
|
||||
struct clip_graph_kimik25 : clip_graph {
|
||||
clip_graph_kimik25(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
@@ -0,0 +1,766 @@
|
||||
#include "models.h"
|
||||
|
||||
#include <string>
|
||||
|
||||
// on-device sampling: top-k, top-p, then a random draw
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::do_sampling(ggml_tensor * logits, ggml_tensor * inp_rand) const {
|
||||
logits = ggml_reshape_1d(ctx0, logits, ggml_nelements(logits));
|
||||
const int64_t n_vocab = logits->ne[0];
|
||||
|
||||
// sort a's rows by idx
|
||||
auto sort_by = [this](ggml_tensor * a, ggml_tensor * idx) {
|
||||
ggml_tensor * a2d = ggml_reshape_2d(ctx0, a, 1, a->ne[0]);
|
||||
return ggml_reshape_1d(ctx0, ggml_get_rows(ctx0, a2d, idx), idx->ne[0]);
|
||||
};
|
||||
|
||||
ggml_tensor * cur = logits;
|
||||
ggml_tensor * candidates = nullptr; // maps row index back to vocab id
|
||||
|
||||
if (top_k > 0 && top_k < n_vocab) {
|
||||
ggml_tensor * idx = ggml_top_k(ctx0, cur, top_k);
|
||||
candidates = idx;
|
||||
cur = sort_by(cur, idx);
|
||||
cb(cur, "sample_top_k_logits", -1);
|
||||
}
|
||||
|
||||
if (top_p < 1.0f) {
|
||||
ggml_tensor * sorted_idx = ggml_argsort(ctx0, cur, GGML_SORT_ORDER_DESC);
|
||||
ggml_tensor * sorted_logits = sort_by(cur, sorted_idx);
|
||||
candidates = candidates ? sort_by(candidates, sorted_idx) : sorted_idx;
|
||||
|
||||
ggml_tensor * probs = ggml_soft_max(ctx0, sorted_logits);
|
||||
ggml_tensor * cdf = ggml_cumsum(ctx0, probs);
|
||||
|
||||
// keep_mask[i] = 1 once cdf[i] crosses top_p
|
||||
ggml_tensor * cdf_scaled = ggml_scale_bias(ctx0, cdf, -1.0f, top_p);
|
||||
ggml_tensor * keep_mask = ggml_step(ctx0, cdf_scaled);
|
||||
ggml_tensor * idxf = ggml_sum(ctx0, keep_mask);
|
||||
idxf = ggml_clamp(ctx0, idxf, 0.0f, (float) keep_mask->ne[0] - 1);
|
||||
ggml_tensor * ones = ggml_scale_bias(ctx0, idxf, 0.0f, 1.0f);
|
||||
|
||||
// top-p must include the crossing element, so force it to 1
|
||||
ggml_tensor * keep_mask_2d = ggml_reshape_2d(ctx0, keep_mask, 1, keep_mask->ne[0]);
|
||||
keep_mask_2d = ggml_set_rows(ctx0, keep_mask_2d, ones, ggml_cast(ctx0, idxf, GGML_TYPE_I32));
|
||||
keep_mask = ggml_reshape_1d(ctx0, keep_mask_2d, keep_mask->ne[0]);
|
||||
|
||||
// log(1) = 0 (keep), log(0) = -inf (drop)
|
||||
ggml_tensor * bias = ggml_log(ctx0, keep_mask);
|
||||
cur = ggml_add(ctx0, sorted_logits, bias);
|
||||
cb(cur, "sample_top_p_logits", -1);
|
||||
}
|
||||
|
||||
// draw one token: find where the cdf crosses inp_rand
|
||||
ggml_tensor * probs = ggml_soft_max(ctx0, cur);
|
||||
ggml_tensor * cumsum = ggml_cumsum(ctx0, probs);
|
||||
|
||||
ggml_tensor * diff = ggml_sub(ctx0, cumsum, inp_rand);
|
||||
ggml_tensor * cross_mask = ggml_step(ctx0, diff);
|
||||
ggml_tensor * idxf = ggml_sum(ctx0, cross_mask);
|
||||
ggml_tensor * idx = ggml_cast(ctx0, ggml_scale_bias(ctx0, idxf, -1.0f, (float) cross_mask->ne[0]), GGML_TYPE_I32);
|
||||
|
||||
if (candidates) {
|
||||
ggml_tensor * cand_2d = ggml_reshape_2d(ctx0, candidates, 1, candidates->ne[0]);
|
||||
idx = ggml_get_rows(ctx0, cand_2d, idx);
|
||||
}
|
||||
cb(idx, "sample_token_id", -1);
|
||||
|
||||
return idx;
|
||||
}
|
||||
|
||||
// returns a new cache with row row_idx set to value
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::cache_set(ggml_tensor * cache, int row_idx, ggml_tensor * value) const {
|
||||
const int64_t n_embd = cache->ne[0];
|
||||
const int64_t n_cache = cache->ne[1];
|
||||
GGML_ASSERT(row_idx >= 0 && row_idx < n_cache);
|
||||
|
||||
// append value as the last row, then gather it back into place
|
||||
ggml_tensor * value_2d = ggml_reshape_2d(ctx0, value, n_embd, 1);
|
||||
ggml_tensor * cache_ext = ggml_concat(ctx0, cache, value_2d, 1); // [n_embd, n_cache + 1]
|
||||
|
||||
// gather indices [0..row_idx-1, n_cache, row_idx+1..n_cache-1]
|
||||
// built via concat, since ggml_set_rows needs F32/F16 values, not an I32 index array
|
||||
ggml_tensor * idx = const_i32(cache, (float) n_cache);
|
||||
if (row_idx > 0) {
|
||||
ggml_tensor * prefix = ggml_cast(ctx0, ggml_arange(ctx0, 0.0f, (float) row_idx, 1.0f), GGML_TYPE_I32);
|
||||
idx = ggml_concat(ctx0, prefix, idx, 0);
|
||||
}
|
||||
if (row_idx < n_cache - 1) {
|
||||
ggml_tensor * suffix = ggml_cast(ctx0, ggml_arange(ctx0, (float) (row_idx + 1), (float) n_cache, 1.0f), GGML_TYPE_I32);
|
||||
idx = ggml_concat(ctx0, idx, suffix, 0);
|
||||
}
|
||||
|
||||
ggml_tensor * result = ggml_get_rows(ctx0, cache_ext, idx);
|
||||
cb(result, "cache_set_out", -1);
|
||||
return result;
|
||||
}
|
||||
|
||||
// builds a const i32 with no host upload: view a tensor, zero it via scale, add value, cast to i32
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::const_i32(ggml_tensor * anchor, float value) const {
|
||||
ggml_tensor * v = ggml_view_1d(ctx0, anchor, 1, 0);
|
||||
if (v->type != GGML_TYPE_F32) {
|
||||
v = ggml_cast(ctx0, v, GGML_TYPE_F32);
|
||||
}
|
||||
return ggml_cast(ctx0, ggml_scale_bias(ctx0, v, 0.0f, value), GGML_TYPE_I32);
|
||||
}
|
||||
|
||||
// causal keep-mask row for a query at position pos, window size n_kv_pad
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::causal_mask_row(int64_t n_kv_pad, int pos) const {
|
||||
ggml_tensor * ones = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv_pad, n_kv_pad), 1.0f);
|
||||
ggml_tensor * keep = ggml_tri(ctx0, ones, GGML_TRI_TYPE_LOWER_DIAG);
|
||||
ggml_tensor * row = ggml_view_1d(ctx0, keep, n_kv_pad, (size_t) pos * keep->nb[1]);
|
||||
ggml_tensor * mask = ggml_log(ctx0, row); // 0 = keep, -inf = masked
|
||||
return ggml_reshape_4d(ctx0, mask, n_kv_pad, 1, 1, 1);
|
||||
}
|
||||
|
||||
// talker hidden size -> predictor hidden size (small_to_mtp_projection)
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::project_in(ggml_tensor * cur) const {
|
||||
if (!model.gen_code_proj_in_w) {
|
||||
return cur;
|
||||
}
|
||||
cur = ggml_mul_mat(ctx0, model.gen_code_proj_in_w, cur);
|
||||
if (model.gen_code_proj_in_b) {
|
||||
cur = ggml_add(ctx0, cur, model.gen_code_proj_in_b);
|
||||
}
|
||||
return cur;
|
||||
}
|
||||
|
||||
// one transformer layer at position pos; writes k/v into k_cache_layer/v_cache_layer at row pos
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::layer_forward(
|
||||
ggml_tensor * cur,
|
||||
const clip_layer & layer,
|
||||
ggml_tensor * inp_pos,
|
||||
ggml_tensor * kq_mask,
|
||||
ggml_tensor *& k_cache_layer,
|
||||
ggml_tensor *& v_cache_layer,
|
||||
int64_t n_kv_pad,
|
||||
int pos,
|
||||
int il) const {
|
||||
const int n_head = hparams.n_head;
|
||||
const int n_head_kv = hparams.n_head_kv;
|
||||
const int64_t d_head = layer.q_w->ne[1] / n_head; // real head_dim, not n_embd / n_head
|
||||
const float kq_scale = 1.0f / sqrtf((float) d_head);
|
||||
|
||||
ggml_tensor * residual = cur;
|
||||
|
||||
ggml_tensor * h = ggml_rms_norm(ctx0, cur, hparams.eps);
|
||||
h = ggml_mul(ctx0, h, layer.ln_1_w);
|
||||
|
||||
ggml_tensor * q = ggml_mul_mat(ctx0, layer.q_w, h);
|
||||
ggml_tensor * k = ggml_mul_mat(ctx0, layer.k_w, h);
|
||||
ggml_tensor * v = ggml_mul_mat(ctx0, layer.v_w, h);
|
||||
|
||||
q = ggml_reshape_3d(ctx0, q, d_head, n_head, 1);
|
||||
k = ggml_reshape_3d(ctx0, k, d_head, n_head_kv, 1);
|
||||
|
||||
q = ggml_rms_norm(ctx0, q, hparams.eps);
|
||||
q = ggml_mul(ctx0, q, layer.q_norm);
|
||||
k = ggml_rms_norm(ctx0, k, hparams.eps);
|
||||
k = ggml_mul(ctx0, k, layer.k_norm);
|
||||
|
||||
q = ggml_rope_ext(ctx0, q, inp_pos, nullptr, (int) d_head, GGML_ROPE_TYPE_NEOX, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
k = ggml_rope_ext(ctx0, k, inp_pos, nullptr, (int) d_head, GGML_ROPE_TYPE_NEOX, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
// write k/v into the cache at row pos, flat layout
|
||||
ggml_tensor * k_flat = ggml_reshape_1d(ctx0, k, d_head * n_head_kv);
|
||||
k_cache_layer = cache_set(k_cache_layer, pos, k_flat);
|
||||
v_cache_layer = cache_set(v_cache_layer, pos, v);
|
||||
|
||||
ggml_tensor * q_cur = ggml_reshape_4d(ctx0, q, d_head, n_head, 1, 1);
|
||||
ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_cache_layer, d_head, n_head_kv, n_kv_pad, 1);
|
||||
ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_cache_layer, d_head, n_head_kv, n_kv_pad, 1);
|
||||
|
||||
ggml_tensor * attn_out = build_attn(layer.o_w, layer.o_b, q_cur, k_cur, v_cur, kq_mask, kq_scale, il);
|
||||
|
||||
cur = ggml_add(ctx0, residual, attn_out);
|
||||
|
||||
ggml_tensor * h2 = ggml_rms_norm(ctx0, cur, hparams.eps);
|
||||
h2 = ggml_mul(ctx0, h2, layer.ln_2_w);
|
||||
|
||||
ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ff_gate_w, h2);
|
||||
ggml_tensor * up = ggml_mul_mat(ctx0, layer.ff_up_w, h2);
|
||||
ggml_tensor * gu = ggml_swiglu_split(ctx0, gate, up);
|
||||
ggml_tensor * down = ggml_mul_mat(ctx0, layer.ff_down_w, gu);
|
||||
|
||||
return ggml_add(ctx0, cur, down);
|
||||
}
|
||||
|
||||
// position 0: hidden bridge, seeds the k/v cache, no sampling
|
||||
// position 1: embed(code0), sample with lm_head[0], write out_code_cache[1]
|
||||
void clip_graph_qwen3tts_gen::code_gen::prefill(
|
||||
std::vector<ggml_tensor *> & k_cache,
|
||||
std::vector<ggml_tensor *> & v_cache,
|
||||
ggml_tensor *& out_code_cache,
|
||||
ggml_tensor * h_state,
|
||||
ggml_tensor * code0_embd,
|
||||
ggml_tensor * inp_rand) const {
|
||||
const int64_t n_kv_pad = k_cache[0]->ne[1];
|
||||
|
||||
{
|
||||
ggml_tensor * cur = project_in(h_state);
|
||||
ggml_tensor * kq_mask = causal_mask_row(n_kv_pad, 0);
|
||||
ggml_tensor * inp_pos = const_i32(k_cache[0], 0.0f);
|
||||
for (size_t il = 0; il < model.layers.size(); il++) {
|
||||
cur = layer_forward(cur, model.layers[il], inp_pos, kq_mask, k_cache[il], v_cache[il], n_kv_pad, 0, (int) il);
|
||||
}
|
||||
// position 0's output is unused, it only seeded the cache
|
||||
}
|
||||
|
||||
{
|
||||
ggml_tensor * cur = project_in(code0_embd);
|
||||
ggml_tensor * kq_mask = causal_mask_row(n_kv_pad, 1);
|
||||
ggml_tensor * inp_pos = const_i32(k_cache[0], 1.0f);
|
||||
for (size_t il = 0; il < model.layers.size(); il++) {
|
||||
cur = layer_forward(cur, model.layers[il], inp_pos, kq_mask, k_cache[il], v_cache[il], n_kv_pad, 1, (int) il);
|
||||
}
|
||||
|
||||
cur = ggml_rms_norm(ctx0, cur, hparams.eps);
|
||||
cur = ggml_mul(ctx0, cur, model.gen_code_norm_w);
|
||||
|
||||
ggml_tensor * head_w = model.gen_code_head_w;
|
||||
ggml_tensor * head_g = ggml_view_2d(ctx0, head_w, head_w->ne[0], head_w->ne[1], head_w->nb[1], 0); // lm_head[0]
|
||||
ggml_tensor * logits = ggml_mul_mat(ctx0, head_g, cur);
|
||||
|
||||
ggml_tensor * sampled = do_sampling(logits, inp_rand);
|
||||
out_code_cache = cache_set(out_code_cache, 1, sampled);
|
||||
}
|
||||
}
|
||||
|
||||
// one decode step of code_predictor
|
||||
// at step_idx g:
|
||||
// - read code from out_code_cache[g], then embed it with codebook table g-1
|
||||
// - write new kv at cache row g+1, sample with lm_head[g]
|
||||
// - write result to out_code_cache[g+1]
|
||||
// step_idx must be in [1, n_acoustic - 1]
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::step(
|
||||
std::vector<ggml_tensor *> & k_cache,
|
||||
std::vector<ggml_tensor *> & v_cache,
|
||||
ggml_tensor * out_code_cache,
|
||||
ggml_tensor * inp_rand,
|
||||
int step_idx) const {
|
||||
const int64_t n_acoustic = model.gen_code_head_w->ne[2];
|
||||
GGML_ASSERT(step_idx >= 1 && step_idx < n_acoustic);
|
||||
GGML_ASSERT(k_cache.size() == model.layers.size());
|
||||
GGML_ASSERT(v_cache.size() == model.layers.size());
|
||||
|
||||
const int64_t n_kv_pad = k_cache[0]->ne[1];
|
||||
const int pos = step_idx + 1; // new cache row and RoPE position
|
||||
|
||||
// embed the previous code via this step's codebook table (rows are already scalars)
|
||||
ggml_tensor * code_in = ggml_view_1d(ctx0, out_code_cache, 1, (size_t) step_idx * out_code_cache->nb[1]);
|
||||
|
||||
ggml_tensor * embd_w = model.gen_code_embd_w; // [n_embd_talker, vocab, n_acoustic]
|
||||
ggml_tensor * embd_g = ggml_view_2d(ctx0, embd_w, embd_w->ne[0], embd_w->ne[1], embd_w->nb[1],
|
||||
(size_t) (step_idx - 1) * embd_w->nb[2]);
|
||||
ggml_tensor * cur = ggml_get_rows(ctx0, embd_g, code_in);
|
||||
cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]);
|
||||
cb(cur, "step_embd_in", step_idx);
|
||||
|
||||
cur = project_in(cur);
|
||||
cb(cur, "step_proj_in", step_idx);
|
||||
|
||||
ggml_tensor * kq_mask = causal_mask_row(n_kv_pad, pos);
|
||||
ggml_tensor * inp_pos = const_i32(k_cache[0], (float) pos);
|
||||
|
||||
for (size_t il = 0; il < model.layers.size(); il++) {
|
||||
cur = layer_forward(cur, model.layers[il], inp_pos, kq_mask, k_cache[il], v_cache[il], n_kv_pad, pos, (int) il);
|
||||
cb(cur, "step_layer_out", (int) il);
|
||||
}
|
||||
|
||||
// final norm, this step's lm_head, sample, write the result
|
||||
cur = ggml_rms_norm(ctx0, cur, hparams.eps);
|
||||
cur = ggml_mul(ctx0, cur, model.gen_code_norm_w);
|
||||
|
||||
ggml_tensor * head_w = model.gen_code_head_w; // [n_embd_pred, vocab, n_acoustic]
|
||||
ggml_tensor * head_g = ggml_view_2d(ctx0, head_w, head_w->ne[0], head_w->ne[1], head_w->nb[1],
|
||||
(size_t) step_idx * head_w->nb[2]);
|
||||
ggml_tensor * logits = ggml_mul_mat(ctx0, head_g, cur);
|
||||
cb(logits, "step_logits", step_idx);
|
||||
|
||||
ggml_tensor * sampled = do_sampling(logits, inp_rand);
|
||||
cb(sampled, "step_sampled", step_idx);
|
||||
|
||||
return cache_set(out_code_cache, pos, sampled);
|
||||
}
|
||||
|
||||
// causal conv1d, stride 1: prepend persisted left-context instead of zero-padding, then a plain conv
|
||||
// x: [T, IC] (T-first). w: [K, IC, OC]. state_name empty means K == 1 (no left-context). returns [T, OC]
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::causal_conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int dilation, const std::string & state_name) const {
|
||||
const int K = (int) w->ne[0];
|
||||
const int pad = (K - 1) * dilation;
|
||||
|
||||
ggml_tensor * x_full = x;
|
||||
if (pad > 0) {
|
||||
ggml_tensor * left = state_in.at(state_name); // [pad, IC]
|
||||
x_full = ggml_concat(ctx0, left, x, 0);
|
||||
}
|
||||
ggml_tensor * y = ggml_conv_1d(ctx0, w, x_full, 1, 0, dilation); // [T, OC, 1]
|
||||
y = ggml_reshape_2d(ctx0, y, y->ne[0], y->ne[1]);
|
||||
if (b) {
|
||||
y = ggml_add(ctx0, y, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
|
||||
}
|
||||
if (pad > 0) {
|
||||
ggml_tensor * new_left = ggml_cont(ctx0, ggml_view_2d(ctx0, x_full, pad, x_full->ne[1], x_full->nb[1],
|
||||
(size_t) (x_full->ne[0] - pad) * x_full->nb[0]));
|
||||
state_out.push_back({state_name, new_left});
|
||||
}
|
||||
return y;
|
||||
}
|
||||
|
||||
// causal depthwise conv1d, stride 1, dilation 1, kernel from w's shape.
|
||||
// x: [T, C]. w: [K, 1, C]. returns [T, C]. see causal_conv1d for the state contract.
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::causal_conv1d_dw(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, const std::string & state_name) const {
|
||||
const int K = (int) w->ne[0];
|
||||
const int pad = K - 1;
|
||||
|
||||
ggml_tensor * x_full = x;
|
||||
if (pad > 0) {
|
||||
ggml_tensor * left = state_in.at(state_name); // [pad, C]
|
||||
x_full = ggml_concat(ctx0, left, x, 0);
|
||||
}
|
||||
ggml_tensor * y = ggml_conv_1d_dw(ctx0, w, x_full, 1, 0, 1); // [T, C, 1]
|
||||
y = ggml_reshape_2d(ctx0, y, y->ne[0], y->ne[1]);
|
||||
if (b) {
|
||||
y = ggml_add(ctx0, y, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
|
||||
}
|
||||
if (pad > 0) {
|
||||
ggml_tensor * new_left = ggml_cont(ctx0, ggml_view_2d(ctx0, x_full, pad, x_full->ne[1], x_full->nb[1],
|
||||
(size_t) (x_full->ne[0] - pad) * x_full->nb[0]));
|
||||
state_out.push_back({state_name, new_left});
|
||||
}
|
||||
return y;
|
||||
}
|
||||
|
||||
// causal ConvTranspose1d, the (kernel - stride) overlap tail is kept as state for the next call
|
||||
// x: [T, IC], w: [K, OC, IC]. state_name empty means K == stride (no overlap). returns [T * stride, OC]
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::causal_conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, const std::string & state_name) const {
|
||||
const int K = (int) w->ne[0];
|
||||
const int OC = (int) w->ne[1];
|
||||
const int trim = K - stride;
|
||||
const int64_t emit_len = x->ne[0] * stride;
|
||||
|
||||
// transposed conv as GEMM + col2im scatter-add, y: [emit_len + trim, OC]
|
||||
ggml_tensor * w2 = ggml_reshape_2d(ctx0, w, (int64_t) K * OC, w->ne[2]);
|
||||
w2 = ggml_cont(ctx0, ggml_transpose(ctx0, w2));
|
||||
ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x));
|
||||
ggml_tensor * col = ggml_mul_mat(ctx0, w2, xt);
|
||||
ggml_tensor * y = ggml_col2im_1d(ctx0, col, stride, OC, 0);
|
||||
|
||||
ggml_tensor * out = y;
|
||||
if (trim > 0) {
|
||||
ggml_tensor * tail = state_in.at(state_name); // [trim, OC]
|
||||
ggml_tensor * head = ggml_add(ctx0, ggml_view_2d(ctx0, y, trim, y->ne[1], y->nb[1], 0), tail);
|
||||
if (emit_len > trim) {
|
||||
ggml_tensor * middle = ggml_view_2d(ctx0, y, emit_len - trim, y->ne[1], y->nb[1], (size_t) trim * y->nb[0]);
|
||||
out = ggml_concat(ctx0, head, middle, 0);
|
||||
} else {
|
||||
out = head;
|
||||
}
|
||||
ggml_tensor * new_tail = ggml_cont(ctx0, ggml_view_2d(ctx0, y, trim, y->ne[1], y->nb[1], (size_t) emit_len * y->nb[0]));
|
||||
state_out.push_back({state_name, new_tail});
|
||||
}
|
||||
if (b) {
|
||||
out = ggml_add(ctx0, out, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// SnakeBeta activation: y = x + sin(alpha*x)^2 * inv_beta (alpha/inv_beta folded via exp/reciprocal at conversion time)
|
||||
// x: [T, C]. alpha/beta: [C], broadcasts over T
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::snake(ggml_tensor * x, ggml_tensor * alpha, ggml_tensor * beta) const {
|
||||
ggml_tensor * a = ggml_reshape_2d(ctx0, alpha, 1, alpha->ne[0]);
|
||||
ggml_tensor * b = ggml_reshape_2d(ctx0, beta, 1, beta->ne[0]);
|
||||
|
||||
// expand reshapes first so mul/sin/sqr/mul/add lands as consecutive nodes, letting backends fuse them
|
||||
ggml_build_forward_expand(gf, a);
|
||||
ggml_build_forward_expand(gf, b);
|
||||
|
||||
ggml_tensor * s = ggml_sin(ctx0, ggml_mul(ctx0, x, a));
|
||||
s = ggml_sqr(ctx0, s);
|
||||
s = ggml_mul(ctx0, s, b);
|
||||
return ggml_add(ctx0, x, s);
|
||||
}
|
||||
|
||||
// RVQ codebook decode: T frames of 16 codes -> 512-dim hidden (C-first, [512, T])
|
||||
// codebook 0 (semantic) and 1..15 (acoustic) sum within their group, project separately, then add
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::quant_decode(ggml_tensor * inp_codes) const {
|
||||
const auto & c2w = model.c2w;
|
||||
const int64_t T = inp_codes->ne[0];
|
||||
|
||||
// ids for codebook group g over all T frames, [T] I32
|
||||
auto group_ids = [&](int g) {
|
||||
return ggml_view_1d(ctx0, inp_codes, T, (size_t) g * inp_codes->nb[1]);
|
||||
};
|
||||
|
||||
ggml_tensor * sem = ggml_get_rows(ctx0, c2w.quant_first_cb_w, group_ids(0)); // [256, T]
|
||||
ggml_tensor * sem_out = ggml_mul_mat(ctx0, c2w.quant_first_out_w, sem); // [512, T]
|
||||
|
||||
ggml_tensor * acc = nullptr;
|
||||
const int64_t n_acoustic = c2w.quant_rest_cb_w->ne[2];
|
||||
for (int g = 1; g <= n_acoustic; g++) {
|
||||
ggml_tensor * cb_g = ggml_view_2d(ctx0, c2w.quant_rest_cb_w, c2w.quant_rest_cb_w->ne[0], c2w.quant_rest_cb_w->ne[1],
|
||||
c2w.quant_rest_cb_w->nb[1], (size_t) (g - 1) * c2w.quant_rest_cb_w->nb[2]);
|
||||
ggml_tensor * embd = ggml_get_rows(ctx0, cb_g, group_ids(g)); // [256, T]
|
||||
acc = acc ? ggml_add(ctx0, acc, embd) : embd;
|
||||
}
|
||||
ggml_tensor * ac_out = ggml_mul_mat(ctx0, c2w.quant_rest_out_w, acc); // [512, T]
|
||||
|
||||
ggml_tensor * hidden = ggml_add(ctx0, sem_out, ac_out);
|
||||
cb(hidden, "wav_quant_hidden", -1);
|
||||
return hidden;
|
||||
}
|
||||
|
||||
// one pre_transformer layer over a batch of N = sliding_window new frames
|
||||
// attention runs over [(W-1)-frame prefix from the last batch] + [N new frames]
|
||||
// RoPE positions come from a persisted counter, so phases line up across batches
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, int il) const {
|
||||
const int n_head = hparams.wav_tfm_n_head;
|
||||
const int n_head_kv = hparams.wav_tfm_n_head_kv;
|
||||
const int64_t d_head = layer.q_w->ne[1] / n_head;
|
||||
const float kq_scale = 1.0f / sqrtf((float) d_head);
|
||||
const int64_t W = hparams.wav_tfm_swa; // == N, frames per batch
|
||||
const int64_t N = cur->ne[1];
|
||||
const int64_t prefix = W - 1;
|
||||
const int64_t total_kv = prefix + N;
|
||||
|
||||
ggml_tensor * residual = cur;
|
||||
ggml_tensor * h = ggml_rms_norm(ctx0, cur, hparams.wav_tfm_eps);
|
||||
h = ggml_mul(ctx0, h, layer.ln_1_w);
|
||||
|
||||
ggml_tensor * q = ggml_mul_mat(ctx0, layer.q_w, h); // [n_head*d_head, N]
|
||||
ggml_tensor * k = ggml_mul_mat(ctx0, layer.k_w, h); // [n_head_kv*d_head, N]
|
||||
ggml_tensor * v = ggml_mul_mat(ctx0, layer.v_w, h); // [n_head_kv*d_head, N]
|
||||
|
||||
q = ggml_reshape_3d(ctx0, q, d_head, n_head, N);
|
||||
k = ggml_reshape_3d(ctx0, k, d_head, n_head_kv, N);
|
||||
|
||||
// real, ever-increasing positions: base (persisted) .. base+N-1
|
||||
ggml_tensor * base = ggml_reshape_1d(ctx0, state_in.at("tfm_pos"), 1);
|
||||
ggml_tensor * offset = ggml_arange(ctx0, 0.0f, (float) N, 1.0f);
|
||||
ggml_tensor * pos = ggml_cast(ctx0, ggml_add(ctx0, offset, base), GGML_TYPE_I32);
|
||||
|
||||
q = ggml_rope_ext(ctx0, q, pos, nullptr, (int) d_head, GGML_ROPE_TYPE_NEOX, 0,
|
||||
hparams.wav_tfm_rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
k = ggml_rope_ext(ctx0, k, pos, nullptr, (int) d_head, GGML_ROPE_TYPE_NEOX, 0,
|
||||
hparams.wav_tfm_rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
// the position counter is the same for all layers, push it once from layer 0
|
||||
if (il == 0) {
|
||||
state_out.push_back({"tfm_pos", ggml_scale_bias(ctx0, state_in.at("tfm_pos"), 1.0f, (float) N)});
|
||||
}
|
||||
|
||||
ggml_tensor * k_new = ggml_reshape_2d(ctx0, k, d_head * n_head_kv, N);
|
||||
ggml_tensor * v_new = ggml_reshape_2d(ctx0, v, d_head * n_head_kv, N);
|
||||
|
||||
ggml_tensor * old_k = state_in.at("tfm_k_" + std::to_string(il)); // [d_head*n_head_kv, W-1]
|
||||
ggml_tensor * old_v = state_in.at("tfm_v_" + std::to_string(il));
|
||||
|
||||
ggml_tensor * k_full = ggml_concat(ctx0, old_k, k_new, 1); // [.., prefix+N]
|
||||
ggml_tensor * v_full = ggml_concat(ctx0, old_v, v_new, 1);
|
||||
|
||||
// next batch's prefix: the last (W-1) frames of this batch
|
||||
state_out.push_back({"tfm_k_" + std::to_string(il),
|
||||
ggml_cont(ctx0, ggml_view_2d(ctx0, k_full, k_full->ne[0], prefix, k_full->nb[1], (size_t) N * k_full->nb[1]))});
|
||||
state_out.push_back({"tfm_v_" + std::to_string(il),
|
||||
ggml_cont(ctx0, ggml_view_2d(ctx0, v_full, v_full->ne[0], prefix, v_full->nb[1], (size_t) N * v_full->nb[1]))});
|
||||
|
||||
// banded causal mask: key j is visible to query i iff 0 <= (prefix+i) - j < W
|
||||
ggml_tensor * pos_k = ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) total_kv, 1.0f), total_kv, 1);
|
||||
ggml_tensor * pos_q = ggml_reshape_2d(ctx0, ggml_arange(ctx0, (float) prefix, (float) (prefix + N), 1.0f), 1, N);
|
||||
ggml_tensor * pos_q_grid = ggml_repeat_4d(ctx0, pos_q, total_kv, N, 1, 1);
|
||||
ggml_tensor * diff = ggml_sub(ctx0, pos_q_grid, pos_k); // [total_kv, N]
|
||||
|
||||
ggml_tensor * causal_keep = ggml_step(ctx0, ggml_scale_bias(ctx0, diff, 1.0f, 0.5f)); // diff >= 0
|
||||
ggml_tensor * in_window = ggml_step(ctx0, ggml_scale_bias(ctx0, diff, -1.0f, (float) W - 0.5f)); // diff < W
|
||||
ggml_tensor * keep = ggml_mul(ctx0, causal_keep, in_window);
|
||||
|
||||
// on a cold start, key j is real state only when j >= prefix - tfm_pos, mask out the rest
|
||||
ggml_tensor * warm = ggml_step(ctx0, ggml_scale_bias(ctx0, ggml_add(ctx0, pos_k, base),
|
||||
1.0f, 0.5f - (float) prefix)); // j + pos > prefix - 0.5
|
||||
keep = ggml_mul(ctx0, keep, warm);
|
||||
|
||||
ggml_tensor * mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, keep), total_kv, N, 1, 1); // 0 = keep, -inf = masked
|
||||
|
||||
ggml_tensor * q_cur = ggml_reshape_4d(ctx0, q, d_head, n_head, N, 1);
|
||||
ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_full, d_head, n_head_kv, total_kv, 1);
|
||||
ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_full, d_head, n_head_kv, total_kv, 1);
|
||||
|
||||
ggml_tensor * attn_out = build_attn(layer.o_w, layer.o_b, q_cur, k_cur, v_cur, mask, kq_scale, il);
|
||||
if (layer.ls_1_w) {
|
||||
attn_out = ggml_mul(ctx0, attn_out, layer.ls_1_w);
|
||||
}
|
||||
cur = ggml_add(ctx0, residual, attn_out);
|
||||
|
||||
ggml_tensor * residual2 = cur;
|
||||
ggml_tensor * h2 = ggml_rms_norm(ctx0, cur, hparams.wav_tfm_eps);
|
||||
h2 = ggml_mul(ctx0, h2, layer.ln_2_w);
|
||||
|
||||
ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ff_gate_w, h2);
|
||||
ggml_tensor * up = ggml_mul_mat(ctx0, layer.ff_up_w, h2);
|
||||
ggml_tensor * gu = ggml_swiglu_split(ctx0, gate, up);
|
||||
ggml_tensor * down = ggml_mul_mat(ctx0, layer.ff_down_w, gu);
|
||||
if (layer.ls_2_w) {
|
||||
down = ggml_mul(ctx0, down, layer.ls_2_w);
|
||||
}
|
||||
return ggml_add(ctx0, residual2, down);
|
||||
}
|
||||
|
||||
// dwconv -> LayerNorm -> pwconv1 -> GELU -> pwconv2 -> layer scale -> residual
|
||||
// x: [T, C] T-first; LayerNorm/pwconv need C on ne0, so this transposes in and back out
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::convnext_block(ggml_tensor * x, const clip_code2wav::upsample_block & blk, const std::string & state_prefix) const {
|
||||
ggml_tensor * residual = x;
|
||||
|
||||
ggml_tensor * h = causal_conv1d_dw(x, blk.dwconv_w, blk.dwconv_b, state_prefix + "_dwconv"); // [T, C]
|
||||
ggml_tensor * hc = ggml_cont(ctx0, ggml_transpose(ctx0, h)); // [C, T]
|
||||
|
||||
hc = ggml_norm(ctx0, hc, 1e-6f);
|
||||
hc = ggml_mul(ctx0, hc, blk.norm_w);
|
||||
hc = ggml_add(ctx0, hc, blk.norm_b);
|
||||
|
||||
ggml_tensor * g = ggml_mul_mat(ctx0, blk.pw1_w, hc);
|
||||
g = ggml_add(ctx0, g, blk.pw1_b);
|
||||
g = ggml_gelu(ctx0, g);
|
||||
g = ggml_mul_mat(ctx0, blk.pw2_w, g);
|
||||
g = ggml_add(ctx0, g, blk.pw2_b);
|
||||
g = ggml_mul(ctx0, g, blk.gamma);
|
||||
|
||||
ggml_tensor * g_t = ggml_cont(ctx0, ggml_transpose(ctx0, g)); // back to [T, C]
|
||||
return ggml_add(ctx0, residual, g_t);
|
||||
}
|
||||
|
||||
// SnakeBeta -> dilated causal conv (k=7) -> SnakeBeta -> pointwise causal conv (k=1) -> residual.
|
||||
// x: [T, C]. returns [T, C].
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::dac_res_unit(ggml_tensor * x, const clip_code2wav::dac_res & res, int dilation, const std::string & state_name) const {
|
||||
ggml_tensor * residual = x;
|
||||
ggml_tensor * h = snake(x, res.act1_alpha, res.act1_beta);
|
||||
h = causal_conv1d(h, res.conv1_w, res.conv1_b, dilation, state_name);
|
||||
h = snake(h, res.act2_alpha, res.act2_beta);
|
||||
h = causal_conv1d(h, res.conv2_w, res.conv2_b, 1, ""); // k=1, no left-context needed
|
||||
return ggml_add(ctx0, residual, h);
|
||||
}
|
||||
|
||||
// RVQ codes -> raw PCM for a batch of N = sliding_window frames
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::decode(ggml_tensor * inp_codes) const {
|
||||
const auto & c2w = model.c2w;
|
||||
|
||||
// 1. quantizer decode: N frames of 16 codes -> [512, N] (C-first)
|
||||
ggml_tensor * hidden = quant_decode(inp_codes);
|
||||
|
||||
// 2. pre_conv: [512, N] -> T-first [N, 512] -> causal conv k=3 -> [N, 1024]
|
||||
ggml_tensor * x = ggml_cont(ctx0, ggml_transpose(ctx0, hidden)); // [N, 512]
|
||||
x = causal_conv1d(x, c2w.pre_conv_w, c2w.pre_conv_b, 1, "pre_conv"); // [N, 1024]
|
||||
cb(x, "wav_pre_conv_out", -1);
|
||||
|
||||
// 3. pre_transformer: back to C-first [1024, N], project down, run the layers, project back up
|
||||
ggml_tensor * cur = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [1024, N]
|
||||
cur = ggml_mul_mat(ctx0, c2w.tfm_in_proj_w, cur);
|
||||
cur = ggml_add(ctx0, cur, c2w.tfm_in_proj_b); // [512 (tfm hidden), N]
|
||||
|
||||
for (int il = 0; il < hparams.wav_tfm_n_layer; il++) {
|
||||
cur = tfm_layer_forward(cur, c2w.tfm_layers[il], il);
|
||||
}
|
||||
|
||||
cur = ggml_rms_norm(ctx0, cur, hparams.wav_tfm_eps);
|
||||
cur = ggml_mul(ctx0, cur, c2w.tfm_output_norm_w);
|
||||
cur = ggml_mul_mat(ctx0, c2w.tfm_out_proj_w, cur);
|
||||
cur = ggml_add(ctx0, cur, c2w.tfm_out_proj_b); // [1024, N]
|
||||
cb(cur, "wav_tfm_out", -1);
|
||||
|
||||
// 4. upsample: 2x (causal ConvTranspose1d, stride 2 + ConvNeXt block), back to T-first
|
||||
// kernel == stride here, so there is no overlap tail to persist
|
||||
x = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [N, 1024]
|
||||
for (size_t il = 0; il < c2w.upsample.size(); il++) {
|
||||
const auto & up = c2w.upsample[il];
|
||||
x = causal_conv_transpose1d(x, up.conv_w, up.conv_b, 2, "");
|
||||
x = convnext_block(x, up, "up" + std::to_string(il));
|
||||
cb(x, "wav_upsample_out", (int) il);
|
||||
}
|
||||
|
||||
// 5. DAC decoder: conv_pre -> n blocks (SnakeBeta -> ConvTranspose1d -> 3 res units) -> conv_post
|
||||
static constexpr int DAC_DILATIONS[3] = { 1, 3, 9 };
|
||||
|
||||
x = causal_conv1d(x, c2w.dac_entry_w, c2w.dac_entry_b, 1, "dac_entry");
|
||||
cb(x, "wav_dac_entry_out", -1);
|
||||
|
||||
for (size_t il = 0; il < c2w.dac.size(); il++) {
|
||||
const auto & blk = c2w.dac[il];
|
||||
const int stride = (int) (blk.conv_w->ne[0] / 2); // kernel == 2*stride for all 4 blocks
|
||||
const std::string blk_name = "dac" + std::to_string(il);
|
||||
x = snake(x, blk.snake_alpha, blk.snake_beta);
|
||||
x = causal_conv_transpose1d(x, blk.conv_w, blk.conv_b, stride, blk_name + "_tail");
|
||||
for (size_t ir = 0; ir < blk.res.size(); ir++) {
|
||||
x = dac_res_unit(x, blk.res[ir], DAC_DILATIONS[ir], blk_name + "_res" + std::to_string(ir));
|
||||
}
|
||||
cb(x, "wav_dac_block_out", (int) il);
|
||||
}
|
||||
|
||||
x = snake(x, c2w.dac_post_snake_alpha, c2w.dac_post_snake_beta);
|
||||
x = causal_conv1d(x, c2w.dac_post_conv_w, c2w.dac_post_conv_b, 1, "dac_post_conv"); // [n_samples, 1]
|
||||
|
||||
x = ggml_clamp(ctx0, x, -1.0f, 1.0f);
|
||||
x = ggml_reshape_1d(ctx0, x, x->ne[0]);
|
||||
cb(x, "wav_audio_out", -1);
|
||||
return x;
|
||||
}
|
||||
|
||||
// code2wav's persisted state buffers: RoPE position counter, K/V per pre_transformer layer,
|
||||
// left-context/tail per stateful conv. shape lookup only, no graph needed
|
||||
std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model) {
|
||||
const auto & c2w = model.c2w;
|
||||
std::vector<c2w_state_slot> slots;
|
||||
|
||||
slots.push_back({"tfm_pos", 1, 1});
|
||||
|
||||
// prefix is (W-1) frames, the batch itself gives the other N=W frames (see tfm_layer_forward)
|
||||
const int64_t d_head = c2w.tfm_layers[0].q_w->ne[1] / hparams.wav_tfm_n_head;
|
||||
const int64_t kv_ch = d_head * hparams.wav_tfm_n_head_kv;
|
||||
const int64_t prefix = hparams.wav_tfm_swa - 1;
|
||||
for (int il = 0; il < hparams.wav_tfm_n_layer; il++) {
|
||||
slots.push_back({"tfm_k_" + std::to_string(il), kv_ch, prefix});
|
||||
slots.push_back({"tfm_v_" + std::to_string(il), kv_ch, prefix});
|
||||
}
|
||||
|
||||
slots.push_back({"pre_conv", c2w.pre_conv_w->ne[0] - 1, c2w.pre_conv_w->ne[1]});
|
||||
|
||||
for (size_t il = 0; il < c2w.upsample.size(); il++) {
|
||||
const auto & up = c2w.upsample[il];
|
||||
slots.push_back({"up" + std::to_string(il) + "_dwconv", up.dwconv_w->ne[0] - 1, up.dwconv_w->ne[2]});
|
||||
}
|
||||
|
||||
slots.push_back({"dac_entry", c2w.dac_entry_w->ne[0] - 1, c2w.dac_entry_w->ne[1]});
|
||||
|
||||
static constexpr int DAC_DILATIONS[3] = { 1, 3, 9 };
|
||||
for (size_t il = 0; il < c2w.dac.size(); il++) {
|
||||
const auto & blk = c2w.dac[il];
|
||||
const int64_t stride = blk.conv_w->ne[0] / 2; // kernel == 2*stride for all 4 blocks
|
||||
const std::string blk_name = "dac" + std::to_string(il);
|
||||
slots.push_back({blk_name + "_tail", stride, blk.conv_w->ne[1]});
|
||||
for (size_t ir = 0; ir < blk.res.size(); ir++) {
|
||||
const auto & res = blk.res[ir];
|
||||
slots.push_back({blk_name + "_res" + std::to_string(ir),
|
||||
(res.conv1_w->ne[0] - 1) * DAC_DILATIONS[ir], res.conv1_w->ne[1]});
|
||||
}
|
||||
}
|
||||
|
||||
slots.push_back({"dac_post_conv", c2w.dac_post_conv_w->ne[0] - 1, c2w.dac_post_conv_w->ne[1]});
|
||||
|
||||
return slots;
|
||||
}
|
||||
|
||||
// both sub-graphs are always built, so the topology stays constant
|
||||
// ggml_build_forward_select() then picks the one that actually runs
|
||||
ggml_cgraph * clip_graph_qwen3tts_gen::build() {
|
||||
GGML_ASSERT(n_batch == 1); // this module only ever processes one frame at a time
|
||||
|
||||
int idx;
|
||||
switch (gen_process) {
|
||||
case CLIP_GEN_PROCESS_GEN_CODE: idx = 0; break;
|
||||
case CLIP_GEN_PROCESS_GEN_WAV: idx = 1; break;
|
||||
default: GGML_ABORT("unknown gen_process");
|
||||
}
|
||||
|
||||
// ---- CLIP_GEN_PROCESS_GEN_CODE: backbone hidden state -> 16 RVQ codes + next-step embd ----
|
||||
// not build_inp_raw(), a GEN_WAV call's `img` has no hidden-state data
|
||||
ggml_tensor * h_state = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_mmproj_embd);
|
||||
ggml_set_name(h_state, "inp_raw"); // must keep this exact name, clip_encode() sets it by name
|
||||
ggml_set_input(h_state);
|
||||
|
||||
ggml_tensor * code0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 1);
|
||||
ggml_set_name(code0, "inp_code0");
|
||||
ggml_set_input(code0);
|
||||
|
||||
ggml_tensor * code0_embd = ggml_get_rows(ctx0, model.gen_code_out_embd_w, code0);
|
||||
code0_embd = ggml_reshape_1d(ctx0, code0_embd, code0_embd->ne[0]);
|
||||
cb(code0_embd, "code0_embd", -1);
|
||||
|
||||
const int64_t n_acoustic = model.gen_code_head_w->ne[2]; // 15
|
||||
const int n_codes = (int) n_acoustic + 1; // 16
|
||||
const int64_t n_kv_pad = n_codes;
|
||||
const int n_layer = (int) model.layers.size();
|
||||
const int n_head = hparams.n_head;
|
||||
const int n_head_kv = hparams.n_head_kv;
|
||||
const int64_t d_head = model.layers[0].q_w->ne[1] / n_head;
|
||||
|
||||
// zero-filled per layer k/v caches, so masked-out rows can't hold garbage
|
||||
std::vector<ggml_tensor *> k_cache(n_layer), v_cache(n_layer);
|
||||
for (int il = 0; il < n_layer; il++) {
|
||||
k_cache[il] = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, d_head * n_head_kv, n_kv_pad), 0.0f);
|
||||
v_cache[il] = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, d_head * n_head_kv, n_kv_pad), 0.0f);
|
||||
}
|
||||
|
||||
code_gen cg(*this, top_k, top_p);
|
||||
|
||||
ggml_tensor * out_code_cache = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_codes);
|
||||
out_code_cache = cg.cache_set(out_code_cache, 0, code0);
|
||||
|
||||
ggml_tensor * inp_rand0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
|
||||
ggml_set_name(inp_rand0, "inp_rand_0");
|
||||
ggml_set_input(inp_rand0);
|
||||
|
||||
cg.prefill(k_cache, v_cache, out_code_cache, h_state, code0_embd, inp_rand0);
|
||||
|
||||
for (int g = 1; g < n_acoustic; g++) {
|
||||
ggml_tensor * inp_rand = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
|
||||
ggml_set_name(inp_rand, ("inp_rand_" + std::to_string(g)).c_str());
|
||||
ggml_set_input(inp_rand);
|
||||
out_code_cache = cg.step(k_cache, v_cache, out_code_cache, inp_rand, g);
|
||||
}
|
||||
|
||||
// output 1: this frame's 16 sampled codes, for the caller's code2wav window
|
||||
ggml_tensor * out_codes = ggml_cont(ctx0, out_code_cache);
|
||||
ggml_set_name(out_codes, "out_codes");
|
||||
ggml_set_output(out_codes);
|
||||
|
||||
// output 2: sum of all 16 codebook embeddings, fed back to the talker for the next frame
|
||||
ggml_tensor * out_embd = code0_embd;
|
||||
for (int g = 1; g <= n_acoustic; g++) {
|
||||
ggml_tensor * code_g = ggml_view_1d(ctx0, out_code_cache, 1, (size_t) g * out_code_cache->nb[1]);
|
||||
|
||||
ggml_tensor * embd_g = ggml_view_2d(ctx0, model.gen_code_embd_w, model.gen_code_embd_w->ne[0], model.gen_code_embd_w->ne[1],
|
||||
model.gen_code_embd_w->nb[1], (size_t) (g - 1) * model.gen_code_embd_w->nb[2]);
|
||||
ggml_tensor * e = ggml_get_rows(ctx0, embd_g, code_g);
|
||||
e = ggml_reshape_1d(ctx0, e, e->ne[0]);
|
||||
|
||||
out_embd = ggml_add(ctx0, out_embd, e);
|
||||
}
|
||||
out_embd = ggml_reshape_2d(ctx0, out_embd, out_embd->ne[0], 1);
|
||||
cb(out_embd, "gen_audio_out", -1);
|
||||
|
||||
// ---- CLIP_GEN_PROCESS_GEN_WAV: 16 RVQ codes -> raw PCM ----
|
||||
const int n_frames = hparams.wav_tfm_swa; // frames per batch, == the attention window
|
||||
|
||||
ggml_tensor * inp_codes = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_frames, n_codes);
|
||||
ggml_set_name(inp_codes, "inp_codes");
|
||||
ggml_set_input(inp_codes);
|
||||
|
||||
code2wav c2w(*this);
|
||||
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
|
||||
ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1);
|
||||
ggml_set_name(t, ("state_in_" + slot.name).c_str());
|
||||
ggml_set_input(t);
|
||||
c2w.state_in[slot.name] = t;
|
||||
}
|
||||
|
||||
ggml_tensor * out_audio = c2w.decode(inp_codes);
|
||||
ggml_set_name(out_audio, "out_audio");
|
||||
ggml_set_output(out_audio);
|
||||
|
||||
for (auto & slot : c2w.state_out) {
|
||||
ggml_set_name(slot.second, ("state_out_" + slot.first).c_str());
|
||||
ggml_set_output(slot.second);
|
||||
}
|
||||
|
||||
// out_embd goes last, clip_encode() reads it back via ggml_graph_node(gf, -1)
|
||||
ggml_tensor * outs[2];
|
||||
outs[0] = out_codes; outs[1] = out_audio;
|
||||
ggml_build_forward_select(gf, outs, 2, idx);
|
||||
for (auto & slot : c2w.state_out) {
|
||||
outs[0] = out_codes; outs[1] = slot.second;
|
||||
ggml_build_forward_select(gf, outs, 2, idx);
|
||||
}
|
||||
outs[0] = out_embd; outs[1] = out_audio;
|
||||
ggml_build_forward_select(gf, outs, 2, idx);
|
||||
|
||||
return gf;
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
#include "models.h"
|
||||
|
||||
static constexpr int SPK_RES2NET_SCALE = 8; // enc_res2net_scale
|
||||
static constexpr int SPK_DILATIONS[3] = { 2, 3, 4 }; // enc_dilations[1..3]
|
||||
|
||||
// conv1d, kernel K, padding "same" (reflect), dilation d
|
||||
// x: [C, T] (ne[0]=C, ne[1]=T) -> [out_c, T]
|
||||
ggml_tensor * clip_graph_qwen3tts_spkenc::conv1d_same(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int dilation) const {
|
||||
const int K = (int) w->ne[0];
|
||||
const int IC = (int) w->ne[1];
|
||||
const int OC = (int) w->ne[2];
|
||||
const int pad = ((K - 1) * dilation) / 2;
|
||||
|
||||
// ggml_pad_reflect_1d pads ne[0], so bring T onto ne[0] first, same layout as im2col wants
|
||||
ggml_tensor * x_t = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T, IC]
|
||||
if (pad > 0) {
|
||||
x_t = ggml_pad_reflect_1d(ctx0, x_t, pad, pad); // [T + 2*pad, IC]
|
||||
}
|
||||
ggml_tensor * x4d = ggml_reshape_4d(ctx0, x_t, x_t->ne[0], IC, 1, 1);
|
||||
|
||||
// dummy F32 kernel, im2col only reads its shape, so a quantized w does not assert
|
||||
ggml_tensor * dummy = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, K, IC, 1, 1);
|
||||
|
||||
ggml_tensor * col = ggml_im2col(ctx0, dummy, x4d, 1, 1, 0, 0, dilation, 1, false, GGML_TYPE_F32);
|
||||
const int64_t T_out = col->ne[1];
|
||||
col = ggml_reshape_2d(ctx0, col, (int64_t) K * IC, T_out);
|
||||
|
||||
ggml_tensor * w2d = ggml_reshape_2d(ctx0, w, (int64_t) K * IC, OC);
|
||||
ggml_tensor * y = ggml_mul_mat(ctx0, w2d, col); // [OC, T_out]
|
||||
ggml_mul_mat_set_prec(y, GGML_PREC_F32);
|
||||
|
||||
ggml_tensor * b2d = ggml_reshape_2d(ctx0, b, OC, 1);
|
||||
y = ggml_add(ctx0, y, b2d);
|
||||
return y;
|
||||
}
|
||||
|
||||
// Res2Net: split channel axis into `scale` chunks, chain dilated conv1d branches
|
||||
// x: [C, T] -> [C, T]
|
||||
ggml_tensor * clip_graph_qwen3tts_spkenc::res2net(ggml_tensor * x, const clip_layer & layer, int dilation, int scale) const {
|
||||
const int64_t C = x->ne[0];
|
||||
const int64_t T = x->ne[1];
|
||||
const int64_t Cs = C / scale;
|
||||
|
||||
std::vector<ggml_tensor *> outs;
|
||||
outs.reserve(scale);
|
||||
|
||||
auto chunk = [&](int i) -> ggml_tensor * {
|
||||
return ggml_view_2d(ctx0, x, Cs, T, x->nb[1], (size_t) i * Cs * x->nb[0]);
|
||||
};
|
||||
|
||||
ggml_tensor * prev = nullptr;
|
||||
for (int i = 0; i < scale; i++) {
|
||||
ggml_tensor * c = ggml_cont(ctx0, chunk(i));
|
||||
if (i == 0) {
|
||||
outs.push_back(c);
|
||||
continue;
|
||||
}
|
||||
ggml_tensor * inp = (i >= 2) ? ggml_add(ctx0, c, prev) : c;
|
||||
ggml_tensor * y = conv1d_same(inp, layer.res2_conv_w[i - 1], layer.res2_conv_b[i - 1], dilation);
|
||||
y = ggml_relu(ctx0, y);
|
||||
outs.push_back(y);
|
||||
prev = y;
|
||||
}
|
||||
|
||||
ggml_tensor * acc = outs[0];
|
||||
for (int i = 1; i < scale; i++) {
|
||||
acc = ggml_concat(ctx0, acc, outs[i], 0);
|
||||
}
|
||||
return acc;
|
||||
}
|
||||
|
||||
// squeeze-and-excitation gate. x: [C, T] -> [C, T]
|
||||
ggml_tensor * clip_graph_qwen3tts_spkenc::se_block(ggml_tensor * x, const clip_layer & layer) const {
|
||||
// temporal mean, keepdim: transpose so T is on ne[0], reduce, transpose back
|
||||
ggml_tensor * x_t = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T, C]
|
||||
ggml_tensor * mean = ggml_mean(ctx0, x_t); // [1, C]
|
||||
mean = ggml_cont(ctx0, ggml_transpose(ctx0, mean)); // [C, 1]
|
||||
|
||||
ggml_tensor * h = conv1d_same(mean, layer.se_conv1_w, layer.se_conv1_b, 1);
|
||||
h = ggml_relu(ctx0, h);
|
||||
h = conv1d_same(h, layer.se_conv2_w, layer.se_conv2_b, 1);
|
||||
h = ggml_sigmoid(ctx0, h); // [C, 1]
|
||||
|
||||
return ggml_mul(ctx0, x, h); // broadcast gate over T
|
||||
}
|
||||
|
||||
// tdnn1 -> res2net -> tdnn2 -> se, plus residual. x: [C, T] -> [C, T]
|
||||
ggml_tensor * clip_graph_qwen3tts_spkenc::se_res2net_block(ggml_tensor * x, const clip_layer & layer, int dilation, int scale) const {
|
||||
ggml_tensor * residual = x;
|
||||
ggml_tensor * h = conv1d_same(x, layer.conv_pw1_w, layer.conv_pw1_b, 1); // tdnn1
|
||||
h = ggml_relu(ctx0, h);
|
||||
h = res2net(h, layer, dilation, scale);
|
||||
h = conv1d_same(h, layer.conv_pw2_w, layer.conv_pw2_b, 1); // tdnn2
|
||||
h = ggml_relu(ctx0, h);
|
||||
h = se_block(h, layer);
|
||||
return ggml_add(ctx0, h, residual);
|
||||
}
|
||||
|
||||
// attentive statistics pooling. x: [C, T] -> [2*C, 1]
|
||||
ggml_tensor * clip_graph_qwen3tts_spkenc::attentive_stats_pool(ggml_tensor * x) const {
|
||||
const int64_t T = x->ne[1];
|
||||
|
||||
// mean over T: [C, 1]
|
||||
ggml_tensor * x_t = ggml_cont(ctx0, ggml_transpose(ctx0, x));
|
||||
ggml_tensor * mean = ggml_mean(ctx0, x_t);
|
||||
mean = ggml_cont(ctx0, ggml_transpose(ctx0, mean));
|
||||
|
||||
// std over T: sqrt(clamp(mean((x - mean)^2), eps))
|
||||
ggml_tensor * mean_rep = ggml_repeat(ctx0, mean, x);
|
||||
ggml_tensor * centered = ggml_sub(ctx0, x, mean_rep);
|
||||
ggml_tensor * var_t = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_sqr(ctx0, centered)));
|
||||
ggml_tensor * var = ggml_mean(ctx0, var_t);
|
||||
var = ggml_cont(ctx0, ggml_transpose(ctx0, var));
|
||||
var = ggml_scale_bias(ctx0, var, 1.0f, 1e-12f);
|
||||
ggml_tensor * std = ggml_sqrt(ctx0, var);
|
||||
|
||||
// attention input: cat([x, mean, std]) along channel axis -> [3C, T]
|
||||
ggml_tensor * std_rep = ggml_repeat(ctx0, std, x);
|
||||
ggml_tensor * cat = ggml_concat(ctx0, x, mean_rep, 0);
|
||||
cat = ggml_concat(ctx0, cat, std_rep, 0);
|
||||
|
||||
// attention TDNN (3C -> attn_c) + ReLU, tanh, then 1x1 conv (attn_c -> C)
|
||||
ggml_tensor * a = conv1d_same(cat, model.spk_asp_tdnn_w, model.spk_asp_tdnn_b, 1);
|
||||
a = ggml_relu(ctx0, a);
|
||||
a = ggml_tanh(ctx0, a);
|
||||
a = conv1d_same(a, model.spk_asp_attn_w, model.spk_asp_attn_b, 1);
|
||||
|
||||
// softmax over T
|
||||
ggml_tensor * a_t = ggml_cont(ctx0, ggml_transpose(ctx0, a)); // [T, C]
|
||||
ggml_tensor * w_t = ggml_soft_max(ctx0, a_t);
|
||||
ggml_tensor * w = ggml_cont(ctx0, ggml_transpose(ctx0, w_t)); // [C, T]
|
||||
|
||||
// weighted mean: sum(w * x) over T, multiply by T to undo ggml_mean's 1/T scaling
|
||||
ggml_tensor * wx = ggml_mul(ctx0, w, x);
|
||||
ggml_tensor * wx_t = ggml_cont(ctx0, ggml_transpose(ctx0, wx));
|
||||
ggml_tensor * w_mean = ggml_mean(ctx0, wx_t);
|
||||
w_mean = ggml_scale(ctx0, w_mean, (float) T);
|
||||
w_mean = ggml_cont(ctx0, ggml_transpose(ctx0, w_mean)); // [C, 1]
|
||||
|
||||
// weighted std: sum(w * (x - w_mean)^2) over T
|
||||
ggml_tensor * w_mean_rep = ggml_repeat(ctx0, w_mean, x);
|
||||
ggml_tensor * dev = ggml_sub(ctx0, x, w_mean_rep);
|
||||
ggml_tensor * w_var_in = ggml_mul(ctx0, w, ggml_sqr(ctx0, dev));
|
||||
ggml_tensor * w_var_t = ggml_cont(ctx0, ggml_transpose(ctx0, w_var_in));
|
||||
ggml_tensor * w_var = ggml_mean(ctx0, w_var_t);
|
||||
w_var = ggml_scale(ctx0, w_var, (float) T);
|
||||
w_var = ggml_cont(ctx0, ggml_transpose(ctx0, w_var));
|
||||
w_var = ggml_scale_bias(ctx0, w_var, 1.0f, 1e-12f);
|
||||
ggml_tensor * w_std = ggml_sqrt(ctx0, w_var);
|
||||
|
||||
return ggml_concat(ctx0, w_mean, w_std, 0); // [2C, 1]
|
||||
}
|
||||
|
||||
ggml_cgraph * clip_graph_qwen3tts_spkenc::build() {
|
||||
// inp_raw: [T, n_mel, 1, 1], from mtmd_audio_preprocessor_qwen3tts_spk
|
||||
ggml_tensor * inp = build_inp_raw(1);
|
||||
inp = ggml_reshape_2d(ctx0, inp, inp->ne[0], inp->ne[1]);
|
||||
|
||||
// this file's convention is [C, T]; the preprocessor delivers [T, C]
|
||||
ggml_tensor * mel = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); // [n_mel, T]
|
||||
cb(mel, "mel", -1);
|
||||
|
||||
// frontend conv0 TDNN k=5, dilation=1: 128 -> 512
|
||||
ggml_tensor * cur = conv1d_same(mel, model.conv1d_1_w, model.conv1d_1_b, 1);
|
||||
cur = ggml_relu(ctx0, cur);
|
||||
cb(cur, "frontend", -1);
|
||||
|
||||
// 3 SE-Res2Net blocks at dilations 2, 3, 4
|
||||
GGML_ASSERT((int) model.layers.size() == 3);
|
||||
std::vector<ggml_tensor *> blk_out(3);
|
||||
for (int il = 0; il < 3; il++) {
|
||||
cur = se_res2net_block(cur, model.layers[il], SPK_DILATIONS[il], SPK_RES2NET_SCALE);
|
||||
blk_out[il] = cur;
|
||||
cb(cur, "block_out", il);
|
||||
}
|
||||
|
||||
// multi-layer feature aggregation: cat blk[0..2] then TDNN k=1 + ReLU
|
||||
ggml_tensor * cat = ggml_concat(ctx0, blk_out[0], blk_out[1], 0);
|
||||
cat = ggml_concat(ctx0, cat, blk_out[2], 0); // [1536, T]
|
||||
ggml_tensor * mfa = conv1d_same(cat, model.conv_out_w, model.conv_out_b, 1);
|
||||
mfa = ggml_relu(ctx0, mfa);
|
||||
cb(mfa, "mfa", -1);
|
||||
|
||||
// attentive statistics pooling: [1536, T] -> [3072, 1]
|
||||
ggml_tensor * stats = attentive_stats_pool(mfa);
|
||||
cb(stats, "asp", -1);
|
||||
|
||||
// final FC k=1: [3072, 1] -> [enc_dim, 1]
|
||||
ggml_tensor * emb = conv1d_same(stats, model.mm_fc_w, model.mm_fc_b, 1);
|
||||
|
||||
emb = ggml_reshape_1d(ctx0, emb, emb->ne[0]);
|
||||
emb = ggml_cont(ctx0, emb);
|
||||
cb(emb, "spk_embedding", -1);
|
||||
|
||||
ggml_build_forward_expand(gf, emb);
|
||||
return gf;
|
||||
}
|
||||
@@ -791,6 +791,66 @@ bool mtmd_audio_preprocessor_mimo_audio::preprocess(const float *
|
||||
return true;
|
||||
}
|
||||
|
||||
//
|
||||
// mtmd_audio_preprocessor_qwen3tts_spk
|
||||
//
|
||||
// same as mel_spectrogram() in modeling_qwen3_tts.py
|
||||
// ECAPA-TDNN takes the whole clip in one pass, so no Whisper-style chunking or normalization
|
||||
//
|
||||
|
||||
void mtmd_audio_preprocessor_qwen3tts_spk::initialize() {
|
||||
cache.fill_sin_cos_table(hparams.audio_n_fft);
|
||||
cache.fill_hann_window(hparams.audio_window_len, true);
|
||||
cache.fill_mel_filterbank_matrix(hparams.n_mel_bins, hparams.audio_n_fft, hparams.audio_sample_rate);
|
||||
}
|
||||
|
||||
bool mtmd_audio_preprocessor_qwen3tts_spk::preprocess(const float * samples,
|
||||
size_t n_samples,
|
||||
std::vector<mtmd_audio_mel> & output) {
|
||||
if (n_samples == 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
GGML_ASSERT(!cache.sin_vals.empty());
|
||||
GGML_ASSERT(!cache.cos_vals.empty());
|
||||
GGML_ASSERT(!cache.filters.data.empty());
|
||||
|
||||
// reflect pad by (n_fft - hop) / 2 = 384, matching center=False STFT framing
|
||||
const int pad = (hparams.audio_n_fft - hparams.audio_hop_len) / 2;
|
||||
if (n_samples < (size_t) pad + 1) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<float> padded(n_samples + 2 * pad, 0.0f);
|
||||
for (int i = 0; i < pad; i++) {
|
||||
padded[i] = samples[pad - i];
|
||||
}
|
||||
std::copy(samples, samples + n_samples, padded.begin() + pad);
|
||||
for (int i = 0; i < pad; i++) {
|
||||
padded[n_samples + pad + i] = samples[n_samples - 2 - i];
|
||||
}
|
||||
|
||||
filter_params params;
|
||||
params.n_mel = hparams.n_mel_bins;
|
||||
params.n_fft_bins = 1 + (hparams.audio_n_fft / 2);
|
||||
params.hann_window_size = hparams.audio_window_len;
|
||||
params.hop_length = hparams.audio_hop_len;
|
||||
params.sample_rate = hparams.audio_sample_rate;
|
||||
params.no_padding = true; // reflect padding already applied above
|
||||
params.use_natural_log = true;
|
||||
params.use_magnitude = true;
|
||||
params.mel_floor = 1e-5f;
|
||||
|
||||
mtmd_audio_mel out;
|
||||
bool ok = log_mel_spectrogram(padded.data(), (int) padded.size(), 4, params, cache, out);
|
||||
if (!ok) {
|
||||
return false;
|
||||
}
|
||||
|
||||
output.push_back(std::move(out));
|
||||
return true;
|
||||
}
|
||||
|
||||
//
|
||||
// mtmd_audio_preprocessor_conformer
|
||||
//
|
||||
|
||||
@@ -120,6 +120,15 @@ struct mtmd_audio_preprocessor_mimo_audio : mtmd_audio_preprocessor {
|
||||
mtmd_audio_cache cache;
|
||||
};
|
||||
|
||||
struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor {
|
||||
mtmd_audio_preprocessor_qwen3tts_spk(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {}
|
||||
void initialize() override;
|
||||
bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) override;
|
||||
|
||||
private:
|
||||
mtmd_audio_cache cache;
|
||||
};
|
||||
|
||||
struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor {
|
||||
mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { }
|
||||
void initialize() override;
|
||||
|
||||
@@ -116,6 +116,14 @@ struct mtmd_cli_context {
|
||||
exit(1);
|
||||
}
|
||||
|
||||
init_vision_context(params);
|
||||
|
||||
if (!mtmd_helper_model_can_chat(lctx, ctx_vision.get())) {
|
||||
LOG_ERR("Model does not support chat mode\n");
|
||||
LOG_ERR("Hint: for TTS models, please use llama-tts\n");
|
||||
exit(1);
|
||||
}
|
||||
|
||||
if (!llama_model_chat_template(model, nullptr) && params.chat_template.empty()) {
|
||||
LOG_ERR("Model does not have chat template.\n");
|
||||
LOG_ERR(" For old llava models, you may need to use '--chat-template vicuna'\n");
|
||||
@@ -129,8 +137,6 @@ struct mtmd_cli_context {
|
||||
chat_history.clear();
|
||||
LOG_INF("%s: chat template example:\n%s\n", __func__, common_chat_format_example(tmpls.get(), params.use_jinja, params.default_template_kwargs).c_str());
|
||||
|
||||
init_vision_context(params);
|
||||
|
||||
// load antiprompt tokens for legacy templates
|
||||
if (params.chat_template == "vicuna") {
|
||||
antiprompt_tokens = common_tokenize(lctx, "ASSISTANT:", false, true);
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
#pragma once
|
||||
|
||||
// shared internal utilities for the mtmd-helper-*.cpp translation units
|
||||
// (mtmd-helper.cpp, mtmd-helper-gen.cpp)
|
||||
// NOT part of the public mtmd-helper.h API
|
||||
|
||||
#include "ggml.h"
|
||||
#include "llama.h"
|
||||
#include "mtmd.h"
|
||||
|
||||
#include <cstdarg>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <vector>
|
||||
|
||||
//
|
||||
// logging
|
||||
//
|
||||
|
||||
struct mtmd_helper_logger {
|
||||
ggml_log_callback default_callback = [](ggml_log_level level, const char * text, void * user_data) {
|
||||
(void) level;
|
||||
(void) user_data;
|
||||
fputs(text, stderr);
|
||||
fflush(stderr);
|
||||
};
|
||||
|
||||
ggml_log_callback log_callback = default_callback;
|
||||
void * log_callback_user_data;
|
||||
|
||||
void log_v(enum ggml_log_level level, const char * format, va_list args) {
|
||||
if (format == NULL) {
|
||||
return;
|
||||
}
|
||||
va_list args_copy;
|
||||
va_copy(args_copy, args);
|
||||
char buffer[128];
|
||||
int len = vsnprintf(buffer, 128, format, args);
|
||||
if (len < 128) {
|
||||
log_callback(level, buffer, log_callback_user_data);
|
||||
} else {
|
||||
char * buffer2 = (char *) calloc(len + 1, sizeof(char));
|
||||
vsnprintf(buffer2, len + 1, format, args_copy);
|
||||
buffer2[len] = 0;
|
||||
log_callback(level, buffer2, log_callback_user_data);
|
||||
free(buffer2);
|
||||
}
|
||||
va_end(args_copy);
|
||||
}
|
||||
|
||||
void log(enum ggml_log_level level, const char * format, ...) {
|
||||
va_list args;
|
||||
va_start(args, format);
|
||||
log_v(level, format, args);
|
||||
va_end(args);
|
||||
}
|
||||
};
|
||||
|
||||
// inline, so all TUs including this header share one instance
|
||||
inline mtmd_helper_logger g_logger;
|
||||
|
||||
#define LOG_DBG(...) g_logger.log(GGML_LOG_LEVEL_DEBUG, __VA_ARGS__)
|
||||
#define LOG_INF(...) g_logger.log(GGML_LOG_LEVEL_INFO, __VA_ARGS__)
|
||||
#define LOG_WRN(...) g_logger.log(GGML_LOG_LEVEL_WARN, __VA_ARGS__)
|
||||
#define LOG_ERR(...) g_logger.log(GGML_LOG_LEVEL_ERROR, __VA_ARGS__)
|
||||
|
||||
//
|
||||
// embd batch
|
||||
//
|
||||
|
||||
// helper struct to make working with embd batch easier
|
||||
// note: this will be removed after llama_batch_ext refactoring
|
||||
struct decode_embd_batch {
|
||||
int n_pos_per_embd;
|
||||
int n_mmproj_embd;
|
||||
std::vector<llama_pos> pos;
|
||||
std::vector<llama_pos> pos_view; // used by mrope
|
||||
std::vector<int32_t> n_seq_id;
|
||||
std::vector<llama_seq_id> seq_id_0;
|
||||
std::vector<llama_seq_id *> seq_ids;
|
||||
std::vector<int8_t> logits;
|
||||
llama_batch batch;
|
||||
decode_embd_batch(float * embd, int32_t n_tokens, int n_pos_per_embd, int n_mmproj_embd) : n_pos_per_embd(n_pos_per_embd), n_mmproj_embd(n_mmproj_embd) {
|
||||
GGML_ASSERT(n_tokens > 0 && n_pos_per_embd > 0 && n_mmproj_embd > 0);
|
||||
pos .resize(n_tokens * n_pos_per_embd);
|
||||
n_seq_id.resize(n_tokens);
|
||||
seq_ids .resize(n_tokens + 1);
|
||||
logits .resize(n_tokens);
|
||||
seq_id_0.resize(1);
|
||||
seq_ids [n_tokens] = nullptr;
|
||||
batch = {
|
||||
/*n_tokens =*/ n_tokens,
|
||||
/*tokens =*/ nullptr,
|
||||
/*embd =*/ embd,
|
||||
/*pos =*/ pos.data(),
|
||||
/*n_seq_id =*/ n_seq_id.data(),
|
||||
/*seq_id =*/ seq_ids.data(),
|
||||
/*logits =*/ logits.data(),
|
||||
};
|
||||
}
|
||||
|
||||
void set_position_normal(llama_pos pos_0, llama_seq_id seq_id) {
|
||||
seq_id_0[0] = seq_id;
|
||||
for (int i = 0; i < batch.n_tokens; i++) {
|
||||
batch.pos [i] = pos_0 + i;
|
||||
batch.n_seq_id[i] = 1;
|
||||
batch.seq_id [i] = seq_id_0.data();
|
||||
batch.logits [i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
// M-RoPE for image
|
||||
void set_position_mrope_2d(const std::vector<mtmd_decoder_pos> & rel_pos, llama_seq_id seq_id) {
|
||||
GGML_ASSERT(n_pos_per_embd == 4);
|
||||
GGML_ASSERT(!rel_pos.empty() && (int32_t)rel_pos.size() == batch.n_tokens);
|
||||
seq_id_0[0] = seq_id;
|
||||
for (int32_t i = 0; i < batch.n_tokens; i++) {
|
||||
pos[i ] = rel_pos[i].t;
|
||||
pos[i + batch.n_tokens ] = rel_pos[i].y;
|
||||
pos[i + batch.n_tokens * 2] = rel_pos[i].x;
|
||||
pos[i + batch.n_tokens * 3] = rel_pos[i].z;
|
||||
}
|
||||
for (int i = 0; i < batch.n_tokens; i++) {
|
||||
batch.n_seq_id[i] = 1;
|
||||
batch.seq_id [i] = seq_id_0.data();
|
||||
batch.logits [i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
// M-RoPE for audio
|
||||
void set_position_mrope_1d(llama_pos pos_0, llama_seq_id seq_id) {
|
||||
GGML_ASSERT(n_pos_per_embd == 4);
|
||||
seq_id_0[0] = seq_id;
|
||||
for (int i = 0; i < batch.n_tokens; i++) {
|
||||
pos[i ] = pos_0 + i;
|
||||
pos[i + batch.n_tokens ] = pos_0 + i;
|
||||
pos[i + batch.n_tokens * 2] = pos_0 + i;
|
||||
pos[i + batch.n_tokens * 3] = pos_0 + i;
|
||||
}
|
||||
for (int i = 0; i < batch.n_tokens; i++) {
|
||||
batch.n_seq_id[i] = 1;
|
||||
batch.seq_id [i] = seq_id_0.data();
|
||||
batch.logits [i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
llama_batch get_view(int offset, int n_tokens) {
|
||||
GGML_ASSERT(offset >= 0 && n_tokens > 0 && offset + n_tokens <= batch.n_tokens);
|
||||
llama_pos * pos_ptr;
|
||||
pos_view.clear();
|
||||
pos_view.reserve(n_tokens * n_pos_per_embd);
|
||||
if (n_pos_per_embd > 1) {
|
||||
// mrope
|
||||
// for example, with layout of src: 1234...1234...1234...1234...
|
||||
// offset 2 will give us dst: 34...34...34...34...
|
||||
for (int i = 0; i < n_pos_per_embd; i++) {
|
||||
// assume n_tokens is less than or equal to batch.n_tokens
|
||||
// batch.n_tokens is number of **total** tokens
|
||||
// n_tokens is number of viewed token
|
||||
size_t src_idx = i * batch.n_tokens + offset;
|
||||
pos_view.insert(pos_view.end(),
|
||||
pos.data() + src_idx,
|
||||
pos.data() + src_idx + n_tokens);
|
||||
}
|
||||
pos_ptr = pos_view.data();
|
||||
} else {
|
||||
// normal
|
||||
pos_ptr = pos.data() + offset;
|
||||
}
|
||||
return {
|
||||
/*n_tokens =*/ n_tokens,
|
||||
/*tokens =*/ nullptr,
|
||||
/*embd =*/ batch.embd + offset * n_mmproj_embd,
|
||||
/*pos =*/ pos_ptr,
|
||||
/*n_seq_id =*/ batch.n_seq_id + offset,
|
||||
/*seq_id =*/ batch.seq_id + offset,
|
||||
/*logits =*/ batch.logits + offset,
|
||||
};
|
||||
}
|
||||
};
|
||||
@@ -0,0 +1,505 @@
|
||||
#include "mtmd.h"
|
||||
#include "mtmd-helper.h"
|
||||
#include "mtmd-helper-common.h"
|
||||
#include "llama.h"
|
||||
#include "../src/llama-ext.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#ifdef MTMD_INTERNAL_HEADER
|
||||
#error "mtmd-helper is a public library outside of mtmd. it must not include internal headers"
|
||||
#endif
|
||||
|
||||
//
|
||||
// Audio generation helpers
|
||||
//
|
||||
|
||||
// --tts-lang codes -> language names used by the codec_language special tokens
|
||||
static const std::unordered_map<std::string, std::string> tts_lang_codes = {
|
||||
{ "zh", "chinese" },
|
||||
{ "en", "english" },
|
||||
{ "de", "german" },
|
||||
{ "it", "italian" },
|
||||
{ "pt", "portuguese" },
|
||||
{ "es", "spanish" },
|
||||
{ "ja", "japanese" },
|
||||
{ "ko", "korean" },
|
||||
{ "fr", "french" },
|
||||
{ "ru", "russian" },
|
||||
};
|
||||
|
||||
static std::string tts_resolve_lang(const std::string & lang) {
|
||||
auto it = tts_lang_codes.find(lang);
|
||||
return it != tts_lang_codes.end() ? it->second : lang;
|
||||
}
|
||||
|
||||
static llama_token find_special_token(const llama_vocab * vocab, const std::string & piece) {
|
||||
const int32_t n = llama_vocab_n_tokens(vocab);
|
||||
for (llama_token t = 0; t < n; t++) {
|
||||
if (piece == llama_vocab_get_text(vocab, t)) {
|
||||
return t;
|
||||
}
|
||||
}
|
||||
return LLAMA_TOKEN_NULL;
|
||||
}
|
||||
|
||||
static bool write_wav16(std::vector<char> & buf, const std::vector<float> & pcm, int32_t rate) {
|
||||
// RIFF chunk sizes are 32-bit; refuse to emit a file with a truncated header
|
||||
if (pcm.size() > ((size_t) UINT32_MAX - 36) / 2) {
|
||||
return false;
|
||||
}
|
||||
const uint32_t data_sz = (uint32_t) (pcm.size() * 2);
|
||||
const uint32_t riff_sz = 36 + data_sz;
|
||||
const uint32_t fmt_sz = 16, byte_rate = (uint32_t) rate * 2;
|
||||
const uint16_t fmt = 1, ch = 1, align = 2, bits = 16;
|
||||
const uint32_t rate32 = (uint32_t) rate;
|
||||
auto put = [&](const void * p, size_t n) {
|
||||
const char * c = (const char *) p;
|
||||
buf.insert(buf.end(), c, c + n);
|
||||
};
|
||||
put("RIFF", 4); put(&riff_sz, 4); put("WAVE", 4);
|
||||
put("fmt ", 4); put(&fmt_sz, 4);
|
||||
put(&fmt, 2); put(&ch, 2); put(&rate32, 4);
|
||||
put(&byte_rate, 4); put(&align, 2); put(&bits, 2);
|
||||
put("data", 4); put(&data_sz, 4);
|
||||
for (float v : pcm) {
|
||||
int16_t s = (int16_t) (std::max(-1.0f, std::min(1.0f, v)) * 32767.0f);
|
||||
put(&s, 2);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
class mtmd_gen_audio_pipeline {
|
||||
public:
|
||||
mtmd_gen_audio_pipeline(llama_context * lctx, mtmd_context * mctx)
|
||||
: lctx(lctx), mctx(mctx), model(llama_get_model(lctx)), vocab(llama_model_get_vocab(model)),
|
||||
n_embd(llama_model_n_embd(model)), info(mtmd_gen_audio_get_info(mctx)) {}
|
||||
virtual ~mtmd_gen_audio_pipeline() = default;
|
||||
|
||||
virtual void reset() = 0;
|
||||
virtual int32_t set_input(const mtmd_helper_gen_audio_inp * inp) = 0;
|
||||
// decodes at most n_batch prompt tokens; returns remaining count (0 = done), <0 on error
|
||||
virtual int32_t step_prompt(int32_t n_batch) = 0;
|
||||
// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token,
|
||||
// those read what they need from h_state_in instead
|
||||
virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) = 0;
|
||||
virtual int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) = 0;
|
||||
|
||||
protected:
|
||||
llama_context * lctx;
|
||||
mtmd_context * mctx;
|
||||
const llama_model * model;
|
||||
const llama_vocab * vocab;
|
||||
int n_embd;
|
||||
mtmd_gen_audio_info info;
|
||||
};
|
||||
|
||||
// Qwen3-TTS: backbone samples codec_0, code_predictor gives the other 15 codebooks,
|
||||
// then code2wav decodes them to PCM
|
||||
class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
|
||||
public:
|
||||
using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline;
|
||||
|
||||
void reset() override {
|
||||
seq_id = 0;
|
||||
pos = 0;
|
||||
codes_buf.clear();
|
||||
c2w_state.clear();
|
||||
audio_pcm.clear();
|
||||
overlay.clear();
|
||||
overlay_idx = 0;
|
||||
h_state_buf.clear();
|
||||
out_buf.clear();
|
||||
prompt_embd_buf.clear();
|
||||
prompt_batch.reset();
|
||||
n_prompt = 0;
|
||||
prompt_pos = 0;
|
||||
}
|
||||
|
||||
int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
|
||||
reset();
|
||||
seq_id = inp->seq_id;
|
||||
|
||||
if (!ensure_cache()) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
const std::string lang = tts_resolve_lang((inp->lang && inp->lang[0]) ? inp->lang : "english");
|
||||
const llama_token c_lang = find_special_token(vocab, ("<|codec_language_" + lang + "|>").c_str());
|
||||
if (c_lang == LLAMA_TOKEN_NULL) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: unknown language '%s'\n", lang.c_str());
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::vector<float> speaker_embd;
|
||||
if (inp->speaker_ref) {
|
||||
if (!encode_speaker(inp->speaker_ref, speaker_embd)) {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
const int n_e = n_embd;
|
||||
auto row = [&](llama_token t) {
|
||||
return std::vector<float>(tok_embd.begin() + (size_t) t * n_e,
|
||||
tok_embd.begin() + (size_t) (t + 1) * n_e);
|
||||
};
|
||||
auto sum_row = [&](llama_token a, llama_token b) {
|
||||
std::vector<float> va = row(a), vb = row(b);
|
||||
for (int i = 0; i < n_e; i++) va[(size_t) i] += vb[(size_t) i];
|
||||
return va;
|
||||
};
|
||||
auto sum_vec = [&](llama_token a, const std::vector<float> & vb) {
|
||||
std::vector<float> va = row(a);
|
||||
for (int i = 0; i < n_e; i++) va[(size_t) i] += vb[(size_t) i];
|
||||
return va;
|
||||
};
|
||||
|
||||
// upstream chat wrap, then slices: [0:3] role, [3:-5] utterance body
|
||||
const std::string full = "<|im_start|>assistant\n" + std::string(inp->prompt, inp->prompt_len) +
|
||||
"<|im_end|>\n<|im_start|>assistant\n";
|
||||
std::vector<llama_token> ids(full.size() + 16);
|
||||
int n_ids = llama_tokenize(vocab, full.c_str(), (int32_t) full.size(), ids.data(), (int32_t) ids.size(),
|
||||
false, true);
|
||||
if (n_ids < 8) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
|
||||
return 1;
|
||||
}
|
||||
ids.resize((size_t) n_ids);
|
||||
|
||||
std::vector<std::vector<float>> prompt;
|
||||
for (int i = 0; i < 3; i++) prompt.push_back(row(ids[(size_t) i]));
|
||||
prompt.push_back(sum_row(tts_pad, c_think));
|
||||
prompt.push_back(sum_row(tts_pad, c_think_b));
|
||||
prompt.push_back(sum_row(tts_pad, c_lang));
|
||||
prompt.push_back(sum_row(tts_pad, c_think_e));
|
||||
if (!speaker_embd.empty()) prompt.push_back(sum_vec(tts_pad, speaker_embd));
|
||||
prompt.push_back(sum_row(tts_bos, codec_pad));
|
||||
for (int i = 3; i < n_ids - 5; i++) prompt.push_back(sum_row(ids[(size_t) i], codec_pad));
|
||||
prompt.push_back(sum_row(tts_eos, codec_pad));
|
||||
prompt.push_back(sum_row(tts_pad, codec_bos));
|
||||
|
||||
n_prompt = (int) prompt.size();
|
||||
|
||||
// the talker uses the qwen3vl interleaved mrope, all sections are equal for a text/codec stream
|
||||
mrope = llama_model_rope_type(model) == LLAMA_ROPE_TYPE_MROPE ||
|
||||
llama_model_rope_type(model) == LLAMA_ROPE_TYPE_IMROPE;
|
||||
const int n_pos_per_embd = mrope ? 4 : 1;
|
||||
|
||||
prompt_embd_buf.resize((size_t) n_prompt * (size_t) n_e);
|
||||
for (int i = 0; i < n_prompt; i++) {
|
||||
memcpy(prompt_embd_buf.data() + (size_t) i * n_e, prompt[(size_t) i].data(), (size_t) n_e * sizeof(float));
|
||||
}
|
||||
|
||||
prompt_batch.reset(new decode_embd_batch(prompt_embd_buf.data(), n_prompt, n_pos_per_embd, n_e));
|
||||
if (mrope) prompt_batch->set_position_mrope_1d(0, seq_id);
|
||||
else prompt_batch->set_position_normal (0, seq_id);
|
||||
prompt_pos = 0;
|
||||
|
||||
pos = 0;
|
||||
top_k = inp->top_k > 0 ? inp->top_k : 50;
|
||||
top_p = inp->top_p > 0 ? inp->top_p : 1.0f;
|
||||
out_type = inp->out_type;
|
||||
|
||||
// the text stream keeps flowing during generation: after frame k, the input adds
|
||||
// trailing text row k on top of the codes embedding, then tts_eos, then tts_pad
|
||||
for (int i = 3; i < n_ids - 5; i++) overlay.push_back(row(ids[(size_t) i]));
|
||||
overlay.push_back(row(tts_eos));
|
||||
overlay.push_back(row(tts_pad));
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t step_prompt(int32_t n_batch) override {
|
||||
GGML_ASSERT(n_batch > 0);
|
||||
if (prompt_pos >= n_prompt) {
|
||||
return 0;
|
||||
}
|
||||
const int32_t n_tokens_batch = std::min(n_batch, n_prompt - prompt_pos);
|
||||
llama_batch batch_view = prompt_batch->get_view(prompt_pos, n_tokens_batch);
|
||||
|
||||
const bool is_last_batch = (prompt_pos + n_tokens_batch) == n_prompt;
|
||||
if (is_last_batch) {
|
||||
batch_view.logits[n_tokens_batch - 1] = 1;
|
||||
}
|
||||
|
||||
if (llama_decode(lctx, batch_view) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n");
|
||||
return -1;
|
||||
}
|
||||
|
||||
pos += n_tokens_batch;
|
||||
prompt_pos += n_tokens_batch;
|
||||
|
||||
if (prompt_pos >= n_prompt) {
|
||||
// prompt fully processed, its embedding buffer is no longer needed
|
||||
prompt_batch.reset();
|
||||
prompt_embd_buf.clear();
|
||||
return 0;
|
||||
}
|
||||
return n_prompt - prompt_pos;
|
||||
}
|
||||
|
||||
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) override {
|
||||
mtmd_gen_inp inp{};
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
|
||||
inp.code0 = sampled - codec_0;
|
||||
inp.embd = const_cast<float *>(h_state_in);
|
||||
inp.top_k = top_k;
|
||||
inp.top_p = top_p;
|
||||
mtmd_gen_out out{};
|
||||
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
codes_buf.insert(codes_buf.end(), out.codes, out.codes + out.n_codes);
|
||||
if (out.n_codes > 0 && codes_buf.size() / out.n_codes >= window_frames) {
|
||||
if (!flush_gen_wav()) {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<float> fb(out.embd, out.embd + n_embd);
|
||||
const auto & ov = overlay[std::min(overlay_idx, overlay.size() - 1)];
|
||||
for (int i = 0; i < n_embd; i++) fb[(size_t) i] += ov[(size_t) i];
|
||||
overlay_idx++;
|
||||
|
||||
const int n_pos_per_embd = mrope ? 4 : 1;
|
||||
decode_embd_batch batch_embd(fb.data(), 1, n_pos_per_embd, n_embd);
|
||||
if (mrope) batch_embd.set_position_mrope_1d(pos, seq_id);
|
||||
else batch_embd.set_position_normal (pos, seq_id);
|
||||
batch_embd.batch.logits[0] = 1;
|
||||
pos++;
|
||||
|
||||
if (llama_decode(lctx, batch_embd.batch) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: decode failed\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
const float * he = llama_get_embeddings_ith(lctx, -1);
|
||||
h_state_buf.assign(he, he + n_embd);
|
||||
*h_state_out = h_state_buf.data();
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
|
||||
if (!flush_gen_wav()) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
*out_sample_rate = info.sample_rate;
|
||||
if (out_n_samples) {
|
||||
*out_n_samples = (int64_t) audio_pcm.size();
|
||||
}
|
||||
|
||||
if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
|
||||
*out_data = (const char *) audio_pcm.data();
|
||||
*out_data_len = audio_pcm.size() * sizeof(float);
|
||||
return 0;
|
||||
}
|
||||
|
||||
out_buf.clear();
|
||||
if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
|
||||
return 1;
|
||||
}
|
||||
*out_data = out_buf.data();
|
||||
*out_data_len = out_buf.size();
|
||||
return 0;
|
||||
}
|
||||
|
||||
private:
|
||||
bool ensure_cache() {
|
||||
if (specials_ok) {
|
||||
return true;
|
||||
}
|
||||
codec_0 = find_special_token(vocab, "<|codec_0|>");
|
||||
codec_bos = find_special_token(vocab, "<|codec_bos|>");
|
||||
codec_eos = find_special_token(vocab, "<|codec_eos_token|>");
|
||||
codec_pad = find_special_token(vocab, "<|codec_pad|>");
|
||||
c_think = find_special_token(vocab, "<|codec_think|>");
|
||||
c_think_b = find_special_token(vocab, "<|codec_think_bos|>");
|
||||
c_think_e = find_special_token(vocab, "<|codec_think_eos|>");
|
||||
tts_pad = find_special_token(vocab, "<tts_pad>");
|
||||
tts_bos = find_special_token(vocab, "<tts_text_bos>");
|
||||
tts_eos = find_special_token(vocab, "<tts_text_eod>");
|
||||
for (llama_token t : { codec_0, codec_bos, codec_eos, codec_pad,
|
||||
c_think, c_think_b, c_think_e,
|
||||
tts_pad, tts_bos, tts_eos }) {
|
||||
if (t == LLAMA_TOKEN_NULL) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: missing a required special token in vocab\n");
|
||||
return false;
|
||||
}
|
||||
}
|
||||
const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr);
|
||||
if (n_tok_embd == 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: model has no token embeddings\n");
|
||||
return false;
|
||||
}
|
||||
tok_embd.resize(n_tok_embd);
|
||||
if (llama_model_get_tok_embd(model, tok_embd.data()) != n_tok_embd) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: token embedding copy failed\n");
|
||||
return false;
|
||||
}
|
||||
specials_ok = true;
|
||||
return true;
|
||||
}
|
||||
|
||||
// runs the reference wav through the speaker encoder, returns one x-vector embedding row
|
||||
bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
|
||||
if (!mtmd_support_audio(mctx)) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n");
|
||||
return false;
|
||||
}
|
||||
const std::string marker = mtmd_default_marker();
|
||||
mtmd_input_text text{ marker.c_str(), marker.size(), false, true };
|
||||
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
|
||||
const mtmd_bitmap * bptr = bitmap;
|
||||
bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
|
||||
if (ok) {
|
||||
ok = false;
|
||||
for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
|
||||
const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
|
||||
if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
|
||||
continue;
|
||||
}
|
||||
if (mtmd_encode_chunk(mctx, chunk) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n");
|
||||
break;
|
||||
}
|
||||
const float * embd = mtmd_get_output_embd(mctx);
|
||||
const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
|
||||
out.assign(embd, embd + n);
|
||||
ok = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
mtmd_input_chunks_free(chunks);
|
||||
return ok;
|
||||
}
|
||||
|
||||
// one GEN_WAV process() call over the buffered codes, state is carried across batches
|
||||
bool flush_gen_wav() {
|
||||
if (codes_buf.empty()) {
|
||||
return true;
|
||||
}
|
||||
mtmd_gen_inp inp{};
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
|
||||
inp.codes = codes_buf.data();
|
||||
inp.n_codes = codes_buf.size();
|
||||
inp.state_data = c2w_state.empty() ? nullptr : (const char *) c2w_state.data();
|
||||
inp.state_size = c2w_state.size();
|
||||
mtmd_gen_out out{};
|
||||
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: gen_wav process failed\n");
|
||||
return false;
|
||||
}
|
||||
audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
|
||||
c2w_state.assign(out.state_data, out.state_data + out.state_size);
|
||||
codes_buf.clear();
|
||||
return true;
|
||||
}
|
||||
|
||||
// vocab specials fixed across the whole session, looked up once
|
||||
bool specials_ok = false;
|
||||
llama_token codec_0 = LLAMA_TOKEN_NULL;
|
||||
llama_token codec_bos = LLAMA_TOKEN_NULL;
|
||||
llama_token codec_eos = LLAMA_TOKEN_NULL;
|
||||
llama_token codec_pad = LLAMA_TOKEN_NULL;
|
||||
llama_token c_think = LLAMA_TOKEN_NULL;
|
||||
llama_token c_think_b = LLAMA_TOKEN_NULL;
|
||||
llama_token c_think_e = LLAMA_TOKEN_NULL;
|
||||
llama_token tts_pad = LLAMA_TOKEN_NULL;
|
||||
llama_token tts_bos = LLAMA_TOKEN_NULL;
|
||||
llama_token tts_eos = LLAMA_TOKEN_NULL;
|
||||
std::vector<float> tok_embd; // whole token embedding matrix, n_vocab * n_embd
|
||||
|
||||
// must match hparams.wav_tfm_swa hardcoded in clip.cpp
|
||||
size_t window_frames = 72;
|
||||
|
||||
// per-generation state, cleared by reset()
|
||||
llama_seq_id seq_id = 0;
|
||||
bool mrope = false;
|
||||
int pos = 0;
|
||||
// prompt decode state, consumed batch-by-batch by step_prompt()
|
||||
std::vector<float> prompt_embd_buf;
|
||||
std::unique_ptr<decode_embd_batch> prompt_batch;
|
||||
int n_prompt = 0;
|
||||
int prompt_pos = 0;
|
||||
int32_t top_k = 50;
|
||||
float top_p = 1.0f;
|
||||
std::vector<int32_t> codes_buf;
|
||||
std::vector<uint8_t> c2w_state;
|
||||
std::vector<float> audio_pcm;
|
||||
std::vector<std::vector<float>> overlay;
|
||||
size_t overlay_idx = 0;
|
||||
std::vector<float> h_state_buf;
|
||||
mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
|
||||
std::vector<char> out_buf;
|
||||
};
|
||||
|
||||
static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lctx, mtmd_context * mctx) {
|
||||
switch (mtmd_gen_audio_get_info(mctx).type) {
|
||||
case MTMD_GEN_AUDIO_TYPE_QWEN3TTS:
|
||||
return std::unique_ptr<mtmd_gen_audio_pipeline>(new qwen3tts_gen_audio_pipeline(lctx, mctx));
|
||||
default:
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
struct mtmd_helper_gen_audio {
|
||||
std::unique_ptr<mtmd_gen_audio_pipeline> pipeline;
|
||||
};
|
||||
|
||||
mtmd_helper_gen_audio * mtmd_helper_gen_audio_init(struct llama_context * lctx, struct mtmd_context * mctx) {
|
||||
auto * ctx = new mtmd_helper_gen_audio();
|
||||
ctx->pipeline = make_pipeline(lctx, mctx);
|
||||
return ctx;
|
||||
}
|
||||
|
||||
void mtmd_helper_gen_audio_free(mtmd_helper_gen_audio * ctx) {
|
||||
delete ctx;
|
||||
}
|
||||
|
||||
void mtmd_helper_gen_audio_reset(mtmd_helper_gen_audio * ctx) {
|
||||
if (ctx->pipeline) {
|
||||
ctx->pipeline->reset();
|
||||
}
|
||||
}
|
||||
|
||||
int32_t mtmd_helper_gen_audio_set_input(mtmd_helper_gen_audio * ctx, const mtmd_helper_gen_audio_inp * inp) {
|
||||
if (!ctx->pipeline) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: unsupported or missing gen-audio pipeline\n");
|
||||
return 1;
|
||||
}
|
||||
return ctx->pipeline->set_input(inp);
|
||||
}
|
||||
|
||||
int32_t mtmd_helper_gen_audio_step_prompt(mtmd_helper_gen_audio * ctx, int32_t n_batch) {
|
||||
if (!ctx->pipeline) {
|
||||
return -1;
|
||||
}
|
||||
return ctx->pipeline->step_prompt(n_batch);
|
||||
}
|
||||
|
||||
int32_t mtmd_helper_gen_audio_step_gen(mtmd_helper_gen_audio * ctx, llama_token sampled,
|
||||
const float * h_state_in, const float ** h_state_out) {
|
||||
if (!ctx->pipeline) {
|
||||
return 1;
|
||||
}
|
||||
return ctx->pipeline->step_gen(sampled, h_state_in, h_state_out);
|
||||
}
|
||||
|
||||
int32_t mtmd_helper_gen_audio_get_output(mtmd_helper_gen_audio * ctx, int32_t * out_sample_rate,
|
||||
const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) {
|
||||
if (!ctx->pipeline) {
|
||||
return 1;
|
||||
}
|
||||
return ctx->pipeline->get_output(out_sample_rate, out_data, out_data_len, out_n_samples);
|
||||
}
|
||||
+16
-155
@@ -9,6 +9,7 @@
|
||||
|
||||
#include "mtmd.h"
|
||||
#include "mtmd-helper.h"
|
||||
#include "mtmd-helper-common.h"
|
||||
#include "llama.h"
|
||||
|
||||
#include <algorithm>
|
||||
@@ -45,50 +46,6 @@
|
||||
// internal logging functions
|
||||
//
|
||||
|
||||
struct mtmd_helper_logger {
|
||||
ggml_log_callback default_callback = [](ggml_log_level level, const char * text, void * user_data) {
|
||||
(void) level;
|
||||
(void) user_data;
|
||||
fputs(text, stderr);
|
||||
fflush(stderr);
|
||||
};
|
||||
|
||||
ggml_log_callback log_callback = default_callback;
|
||||
void * log_callback_user_data;
|
||||
|
||||
void log_v(enum ggml_log_level level, const char * format, va_list args) {
|
||||
if (format == NULL) {
|
||||
return;
|
||||
}
|
||||
va_list args_copy;
|
||||
va_copy(args_copy, args);
|
||||
char buffer[128];
|
||||
int len = vsnprintf(buffer, 128, format, args);
|
||||
if (len < 128) {
|
||||
log_callback(level, buffer, log_callback_user_data);
|
||||
} else {
|
||||
char * buffer2 = (char *) calloc(len + 1, sizeof(char));
|
||||
vsnprintf(buffer2, len + 1, format, args_copy);
|
||||
buffer2[len] = 0;
|
||||
log_callback(level, buffer2, log_callback_user_data);
|
||||
free(buffer2);
|
||||
}
|
||||
va_end(args_copy);
|
||||
}
|
||||
|
||||
void log(enum ggml_log_level level, const char * format, ...) {
|
||||
va_list args;
|
||||
va_start(args, format);
|
||||
log_v(level, format, args);
|
||||
va_end(args);
|
||||
}
|
||||
} g_logger;
|
||||
|
||||
#define LOG_DBG(...) g_logger.log(GGML_LOG_LEVEL_DEBUG, __VA_ARGS__)
|
||||
#define LOG_INF(...) g_logger.log(GGML_LOG_LEVEL_INFO, __VA_ARGS__)
|
||||
#define LOG_WRN(...) g_logger.log(GGML_LOG_LEVEL_WARN, __VA_ARGS__)
|
||||
#define LOG_ERR(...) g_logger.log(GGML_LOG_LEVEL_ERROR, __VA_ARGS__)
|
||||
|
||||
void mtmd_helper_log_set(ggml_log_callback log_callback, void * user_data) {
|
||||
if (log_callback == nullptr) {
|
||||
log_callback = g_logger.default_callback;
|
||||
@@ -127,117 +84,6 @@ void mtmd_helper_image_get_decoder_pos(const mtmd_image_tokens * chunks, llama_p
|
||||
}
|
||||
}
|
||||
|
||||
// helper struct to make working with embd batch easier
|
||||
// note: this will be removed after llama_batch_ext refactoring
|
||||
struct decode_embd_batch {
|
||||
int n_pos_per_embd;
|
||||
int n_mmproj_embd;
|
||||
std::vector<llama_pos> pos;
|
||||
std::vector<llama_pos> pos_view; // used by mrope
|
||||
std::vector<int32_t> n_seq_id;
|
||||
std::vector<llama_seq_id> seq_id_0;
|
||||
std::vector<llama_seq_id *> seq_ids;
|
||||
std::vector<int8_t> logits;
|
||||
llama_batch batch;
|
||||
decode_embd_batch(float * embd, int32_t n_tokens, int n_pos_per_embd, int n_mmproj_embd) : n_pos_per_embd(n_pos_per_embd), n_mmproj_embd(n_mmproj_embd) {
|
||||
GGML_ASSERT(n_tokens > 0 && n_pos_per_embd > 0 && n_mmproj_embd > 0);
|
||||
pos .resize(n_tokens * n_pos_per_embd);
|
||||
n_seq_id.resize(n_tokens);
|
||||
seq_ids .resize(n_tokens + 1);
|
||||
logits .resize(n_tokens);
|
||||
seq_id_0.resize(1);
|
||||
seq_ids [n_tokens] = nullptr;
|
||||
batch = {
|
||||
/*n_tokens =*/ n_tokens,
|
||||
/*tokens =*/ nullptr,
|
||||
/*embd =*/ embd,
|
||||
/*pos =*/ pos.data(),
|
||||
/*n_seq_id =*/ n_seq_id.data(),
|
||||
/*seq_id =*/ seq_ids.data(),
|
||||
/*logits =*/ logits.data(),
|
||||
};
|
||||
}
|
||||
|
||||
void set_position_normal(llama_pos pos_0, llama_seq_id seq_id) {
|
||||
seq_id_0[0] = seq_id;
|
||||
for (int i = 0; i < batch.n_tokens; i++) {
|
||||
batch.pos [i] = pos_0 + i;
|
||||
batch.n_seq_id[i] = 1;
|
||||
batch.seq_id [i] = seq_id_0.data();
|
||||
batch.logits [i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
// M-RoPE for image
|
||||
void set_position_mrope_2d(const std::vector<mtmd_decoder_pos> & rel_pos, llama_seq_id seq_id) {
|
||||
GGML_ASSERT(n_pos_per_embd == 4);
|
||||
GGML_ASSERT(!rel_pos.empty() && (int32_t)rel_pos.size() == batch.n_tokens);
|
||||
seq_id_0[0] = seq_id;
|
||||
for (int32_t i = 0; i < batch.n_tokens; i++) {
|
||||
pos[i ] = rel_pos[i].t;
|
||||
pos[i + batch.n_tokens ] = rel_pos[i].y;
|
||||
pos[i + batch.n_tokens * 2] = rel_pos[i].x;
|
||||
pos[i + batch.n_tokens * 3] = rel_pos[i].z;
|
||||
}
|
||||
for (int i = 0; i < batch.n_tokens; i++) {
|
||||
batch.n_seq_id[i] = 1;
|
||||
batch.seq_id [i] = seq_id_0.data();
|
||||
batch.logits [i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
// M-RoPE for audio
|
||||
void set_position_mrope_1d(llama_pos pos_0, llama_seq_id seq_id) {
|
||||
GGML_ASSERT(n_pos_per_embd == 4);
|
||||
seq_id_0[0] = seq_id;
|
||||
for (int i = 0; i < batch.n_tokens; i++) {
|
||||
pos[i ] = pos_0 + i;
|
||||
pos[i + batch.n_tokens ] = pos_0 + i;
|
||||
pos[i + batch.n_tokens * 2] = pos_0 + i;
|
||||
pos[i + batch.n_tokens * 3] = pos_0 + i;
|
||||
}
|
||||
for (int i = 0; i < batch.n_tokens; i++) {
|
||||
batch.n_seq_id[i] = 1;
|
||||
batch.seq_id [i] = seq_id_0.data();
|
||||
batch.logits [i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
llama_batch get_view(int offset, int n_tokens) {
|
||||
GGML_ASSERT(offset >= 0 && n_tokens > 0 && offset + n_tokens <= batch.n_tokens);
|
||||
llama_pos * pos_ptr;
|
||||
pos_view.clear();
|
||||
pos_view.reserve(n_tokens * n_pos_per_embd);
|
||||
if (n_pos_per_embd > 1) {
|
||||
// mrope
|
||||
// for example, with layout of src: 1234...1234...1234...1234...
|
||||
// offset 2 will give us dst: 34...34...34...34...
|
||||
for (int i = 0; i < n_pos_per_embd; i++) {
|
||||
// assume n_tokens is less than or equal to batch.n_tokens
|
||||
// batch.n_tokens is number of **total** tokens
|
||||
// n_tokens is number of viewed token
|
||||
size_t src_idx = i * batch.n_tokens + offset;
|
||||
pos_view.insert(pos_view.end(),
|
||||
pos.data() + src_idx,
|
||||
pos.data() + src_idx + n_tokens);
|
||||
}
|
||||
pos_ptr = pos_view.data();
|
||||
} else {
|
||||
// normal
|
||||
pos_ptr = pos.data() + offset;
|
||||
}
|
||||
return {
|
||||
/*n_tokens =*/ n_tokens,
|
||||
/*tokens =*/ nullptr,
|
||||
/*embd =*/ batch.embd + offset * n_mmproj_embd,
|
||||
/*pos =*/ pos_ptr,
|
||||
/*n_seq_id =*/ batch.n_seq_id + offset,
|
||||
/*seq_id =*/ batch.seq_id + offset,
|
||||
/*logits =*/ batch.logits + offset,
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
// Helper class to set non-causal attention via RAII
|
||||
class scope_non_causal {
|
||||
public:
|
||||
@@ -1084,3 +930,18 @@ int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx,
|
||||
GGML_ASSERT(false && "video is not supported in this build (MTMD_VIDEO is set to OFF)");
|
||||
#endif
|
||||
}
|
||||
|
||||
bool mtmd_helper_model_can_chat(llama_context * lctx, mtmd_context * mctx) {
|
||||
if (!mctx) {
|
||||
return true;
|
||||
}
|
||||
|
||||
auto * model = llama_get_model(lctx);
|
||||
auto * tmpl = llama_model_chat_template(model, nullptr);
|
||||
auto info = mtmd_gen_audio_get_info(mctx);
|
||||
|
||||
// tts-only model cannot be used for chat (no chat template)
|
||||
bool is_tts_only = info.type != MTMD_GEN_AUDIO_TYPE_NONE && tmpl == nullptr;
|
||||
|
||||
return !is_tts_only;
|
||||
}
|
||||
|
||||
@@ -157,6 +157,73 @@ MTMD_API int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx,
|
||||
mtmd_bitmap ** out_bitmap,
|
||||
char ** out_text);
|
||||
|
||||
// return true if model can be used for chat
|
||||
MTMD_API bool mtmd_helper_model_can_chat(struct llama_context * lctx, struct mtmd_context * mctx);
|
||||
|
||||
//
|
||||
// Audio generation helpers
|
||||
// (early-stage experimental, subjected to breaking changes)
|
||||
//
|
||||
|
||||
// audio generation helper context
|
||||
// contains accumulator for generated audio features and PCM audio
|
||||
struct mtmd_helper_gen_audio;
|
||||
typedef struct mtmd_helper_gen_audio mtmd_helper_gen_audio;
|
||||
|
||||
enum mtmd_helper_gen_audio_outtype {
|
||||
MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM, // raw PCM
|
||||
MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV, // WAV PCM 16-bit LE, mono
|
||||
};
|
||||
struct mtmd_helper_gen_audio_inp {
|
||||
llama_seq_id seq_id;
|
||||
|
||||
const char * prompt;
|
||||
size_t prompt_len;
|
||||
|
||||
mtmd_bitmap * speaker_ref; // optional, can be NULL
|
||||
const char * lang; // optional, can be NULL
|
||||
|
||||
int32_t top_k;
|
||||
float top_p;
|
||||
|
||||
enum mtmd_helper_gen_audio_outtype out_type;
|
||||
};
|
||||
|
||||
MTMD_API mtmd_helper_gen_audio * mtmd_helper_gen_audio_init(
|
||||
struct llama_context * lctx,
|
||||
struct mtmd_context * mctx);
|
||||
|
||||
MTMD_API void mtmd_helper_gen_audio_free(mtmd_helper_gen_audio * ctx);
|
||||
|
||||
MTMD_API void mtmd_helper_gen_audio_reset(mtmd_helper_gen_audio * ctx);
|
||||
|
||||
MTMD_API int32_t mtmd_helper_gen_audio_set_input(
|
||||
mtmd_helper_gen_audio * ctx,
|
||||
const struct mtmd_helper_gen_audio_inp * inp);
|
||||
|
||||
// processes at most n_batch prompt tokens per call
|
||||
// returns: >0 = number of prompt tokens remaining, 0 = done, <0 = error
|
||||
MTMD_API int32_t mtmd_helper_gen_audio_step_prompt(
|
||||
mtmd_helper_gen_audio * ctx,
|
||||
int32_t n_batch);
|
||||
|
||||
// generates one frame; must only be called after step_prompt() has returned 0
|
||||
// h_state_out is valid until next step_gen() or reset() call
|
||||
MTMD_API int32_t mtmd_helper_gen_audio_step_gen(
|
||||
mtmd_helper_gen_audio * ctx,
|
||||
llama_token sampled,
|
||||
const float * h_state_in,
|
||||
const float ** h_state_out);
|
||||
|
||||
// out_data valid until next get_output() or reset() call
|
||||
// out_n_samples (optional, can be NULL) receives the number of generated PCM samples
|
||||
MTMD_API int32_t mtmd_helper_gen_audio_get_output(
|
||||
mtmd_helper_gen_audio * ctx,
|
||||
int32_t * out_sample_rate,
|
||||
const char ** out_data,
|
||||
size_t * out_data_len,
|
||||
int64_t * out_n_samples);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // extern "C"
|
||||
#endif
|
||||
@@ -177,6 +244,31 @@ struct mtmd_helper_video_deleter {
|
||||
};
|
||||
using video_ptr = std::unique_ptr<mtmd_helper_video, mtmd_helper_video_deleter>;
|
||||
|
||||
// audio generation-related C++ wrappers
|
||||
struct mtmd_helper_gen_audio_deleter {
|
||||
void operator()(mtmd_helper_gen_audio * val) { mtmd_helper_gen_audio_free(val); }
|
||||
};
|
||||
using gen_audio_ptr = std::unique_ptr<mtmd_helper_gen_audio, mtmd_helper_gen_audio_deleter>;
|
||||
struct gen_audio {
|
||||
gen_audio_ptr ctx;
|
||||
gen_audio(struct llama_context * lctx, struct mtmd_context * mctx) : ctx(mtmd_helper_gen_audio_init(lctx, mctx)) {}
|
||||
void reset() {
|
||||
mtmd_helper_gen_audio_reset(ctx.get());
|
||||
}
|
||||
int32_t set_input(const struct mtmd_helper_gen_audio_inp * inp) {
|
||||
return mtmd_helper_gen_audio_set_input(ctx.get(), inp);
|
||||
}
|
||||
int32_t step_prompt(int32_t n_batch) {
|
||||
return mtmd_helper_gen_audio_step_prompt(ctx.get(), n_batch);
|
||||
}
|
||||
int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out) {
|
||||
return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out);
|
||||
}
|
||||
int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples = nullptr) {
|
||||
return mtmd_helper_gen_audio_get_output(ctx.get(), out_sample_rate, out_data, out_data_len, out_n_samples);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace mtmd_helper
|
||||
#endif
|
||||
|
||||
|
||||
@@ -262,6 +262,13 @@ struct mtmd_context {
|
||||
struct clip_ctx * ctx_a; // audio
|
||||
std::vector<float> out_embd; // image embedding vector
|
||||
|
||||
// generation context
|
||||
struct clip_ctx * ctx_gen_a; // audio
|
||||
std::vector<int32_t> gen_out_codes; // this frame's 16 sampled codes (GEN_CODE)
|
||||
std::vector<float> gen_out_embd; // next-step hidden state fed back to backbone (GEN_CODE)
|
||||
std::vector<float> gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV)
|
||||
std::vector<uint8_t> gen_out_state; // state to feed into the next GEN_WAV call
|
||||
|
||||
bool print_timings;
|
||||
int n_threads;
|
||||
std::string media_marker;
|
||||
@@ -354,6 +361,7 @@ struct mtmd_context {
|
||||
auto res = clip_init(mmproj_fname, ctx_clip_params);
|
||||
ctx_v = res.ctx_v;
|
||||
ctx_a = res.ctx_a;
|
||||
ctx_gen_a = res.ctx_gen_a;
|
||||
if (!ctx_v && !ctx_a) {
|
||||
throw std::runtime_error(string_format("Failed to load CLIP model from %s\n", mmproj_fname));
|
||||
}
|
||||
@@ -378,6 +386,15 @@ struct mtmd_context {
|
||||
"hint: you may be using wrong mmproj\n",
|
||||
n_embd_text, n_embd_clip));
|
||||
}
|
||||
if (ctx_gen_a) {
|
||||
int n_embd_gen = clip_n_mmproj_embd(ctx_gen_a);
|
||||
if (n_embd_text > 0 && n_embd_text != n_embd_gen) {
|
||||
throw std::runtime_error(string_format(
|
||||
"mismatch between text model (n_embd = %d) and gen-audio mmproj (n_embd = %d)\n"
|
||||
"hint: you may be using wrong mmproj\n",
|
||||
n_embd_text, n_embd_gen));
|
||||
}
|
||||
}
|
||||
if (ctx_v) {
|
||||
init_vision();
|
||||
}
|
||||
@@ -740,6 +757,10 @@ struct mtmd_context {
|
||||
aud_end = "<|mimo_audio_end|>";
|
||||
audio_preproc = std::make_unique<mtmd_audio_preprocessor_mimo_audio>(ctx_a);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_SPKENC:
|
||||
{
|
||||
audio_preproc = std::make_unique<mtmd_audio_preprocessor_qwen3tts_spk>(ctx_a);
|
||||
} break;
|
||||
default:
|
||||
throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj));
|
||||
}
|
||||
@@ -780,6 +801,7 @@ struct mtmd_context {
|
||||
~mtmd_context() {
|
||||
clip_free(ctx_a);
|
||||
clip_free(ctx_v);
|
||||
clip_free(ctx_gen_a);
|
||||
}
|
||||
|
||||
private:
|
||||
@@ -1553,6 +1575,125 @@ float * mtmd_get_output_embd(mtmd_context * ctx) {
|
||||
return ctx->out_embd.data();
|
||||
}
|
||||
|
||||
//
|
||||
// audio generation
|
||||
//
|
||||
|
||||
mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
|
||||
mtmd_gen_audio_info info;
|
||||
if (!ctx->ctx_gen_a) {
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
|
||||
return info;
|
||||
}
|
||||
switch (clip_get_projector_type(ctx->ctx_gen_a)) {
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_QWEN3TTS;
|
||||
info.sample_rate = 24000;
|
||||
break;
|
||||
default:
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
|
||||
break;
|
||||
}
|
||||
return info;
|
||||
}
|
||||
|
||||
static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
|
||||
clip_ctx * ctx_clip = ctx->ctx_gen_a;
|
||||
if (!ctx_clip) {
|
||||
LOG_ERR("%s: model does not support audio generation\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) {
|
||||
const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip);
|
||||
|
||||
clip_image_f32 hidden_state;
|
||||
hidden_state.set_size({(int) n_embd, 1}, false, true);
|
||||
hidden_state.cpy_buf(std::vector<float>(inp->embd, inp->embd + n_embd));
|
||||
|
||||
clip_image_f32_batch batch;
|
||||
batch.is_audio = true;
|
||||
batch.entries.push_back(std::move(hidden_state));
|
||||
|
||||
std::vector<float> out_embd(n_embd);
|
||||
std::vector<int32_t> out_codes;
|
||||
|
||||
clip_encode_params params;
|
||||
params.imgs = &batch;
|
||||
params.n_threads = ctx->n_threads;
|
||||
params.gen_process = CLIP_GEN_PROCESS_GEN_CODE;
|
||||
params.out_embd = &out_embd;
|
||||
params.out_codes = &out_codes;
|
||||
params.code0 = inp->code0;
|
||||
params.top_k = inp->top_k;
|
||||
params.top_p = inp->top_p;
|
||||
|
||||
if (!clip_encode(ctx_clip, ¶ms)) {
|
||||
LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
ctx->gen_out_embd = std::move(out_embd);
|
||||
ctx->gen_out_codes = std::move(out_codes);
|
||||
|
||||
out->embd = ctx->gen_out_embd.data();
|
||||
out->codes = ctx->gen_out_codes.data();
|
||||
out->n_codes = ctx->gen_out_codes.size();
|
||||
return 0;
|
||||
}
|
||||
|
||||
// MTMD_GEN_PROCESS_TYPE_GEN_WAV
|
||||
if (!inp->codes || inp->n_codes == 0) {
|
||||
LOG_ERR("%s: codes required for gen_wav\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
std::vector<int32_t> in_codes(inp->codes, inp->codes + inp->n_codes);
|
||||
std::vector<uint8_t> in_state;
|
||||
if (inp->state_data) {
|
||||
in_state.assign(inp->state_data, inp->state_data + inp->state_size);
|
||||
}
|
||||
|
||||
// gen_wav has no hidden-state input, the batch entry is an unused placeholder
|
||||
// TODO @ngxson : some models in the future may require hidden-state input, need to update this code later
|
||||
clip_image_f32 dummy;
|
||||
dummy.set_size({1, 1}, false, true);
|
||||
dummy.cpy_buf(std::vector<float>(1, 0.0f));
|
||||
|
||||
clip_image_f32_batch batch;
|
||||
batch.is_audio = true;
|
||||
batch.entries.push_back(std::move(dummy));
|
||||
|
||||
clip_encode_params params;
|
||||
params.imgs = &batch;
|
||||
params.n_threads = ctx->n_threads;
|
||||
params.gen_process = CLIP_GEN_PROCESS_GEN_WAV;
|
||||
params.codes = &in_codes;
|
||||
params.out_audio = &ctx->gen_out_audio;
|
||||
params.state_in = inp->state_data ? &in_state : nullptr;
|
||||
params.state_out = &ctx->gen_out_state;
|
||||
|
||||
if (!clip_encode(ctx_clip, ¶ms)) {
|
||||
LOG_ERR("%s: clip_encode failed (code2wav)\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
out->audio = ctx->gen_out_audio.data();
|
||||
out->n_samples = ctx->gen_out_audio.size();
|
||||
out->state_data = (const char *) ctx->gen_out_state.data();
|
||||
out->state_size = ctx->gen_out_state.size();
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t mtmd_gen_audio_process(mtmd_context * ctx, const struct mtmd_gen_inp * inp, struct mtmd_gen_out * out) {
|
||||
try {
|
||||
return mtmd_gen_audio_process_impl(ctx, inp, out);
|
||||
} catch (const std::exception & e) {
|
||||
LOG_ERR("%s: error: %s\n", __func__, e.what());
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
mtmd_batch * mtmd_batch_init(mtmd_context * ctx) {
|
||||
return new mtmd_batch(ctx);
|
||||
}
|
||||
|
||||
@@ -327,6 +327,60 @@ struct mtmd_caps {
|
||||
};
|
||||
MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname);
|
||||
|
||||
/////////////////////////////////////////
|
||||
// EXPERIMENTAL API for audio generation, subjected to breaking changes
|
||||
|
||||
// represent the pipeline type
|
||||
enum mtmd_gen_audio_type {
|
||||
MTMD_GEN_AUDIO_TYPE_NONE, // not supported
|
||||
MTMD_GEN_AUDIO_TYPE_QWEN3TTS,
|
||||
};
|
||||
struct mtmd_gen_audio_info {
|
||||
enum mtmd_gen_audio_type type;
|
||||
int32_t sample_rate; // in Hz, for example 24000 for qwen3tts
|
||||
};
|
||||
MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx);
|
||||
|
||||
enum mtmd_gen_process_type {
|
||||
MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to semantic (codes, mel-spectrogram, etc.)
|
||||
MTMD_GEN_PROCESS_TYPE_GEN_WAV, // convert semantic to PCM audio
|
||||
// for qwen3tts, this is code2wav
|
||||
};
|
||||
struct mtmd_gen_inp {
|
||||
enum mtmd_gen_process_type type;
|
||||
|
||||
// for MTMD_GEN_PROCESS_TYPE_GEN_CODE
|
||||
int32_t code0; // the sampled codebook 0 entry from backbone
|
||||
float * embd; // the hidden state from backbone, must have n_text_embd elements
|
||||
int32_t top_k;
|
||||
float top_p;
|
||||
|
||||
// for MTMD_GEN_PROCESS_TYPE_GEN_WAV
|
||||
int32_t * codes;
|
||||
size_t n_codes;
|
||||
const char * state_data;
|
||||
size_t state_size;
|
||||
};
|
||||
struct mtmd_gen_out {
|
||||
// note: output memory is allocated by the context, valid until next process() call
|
||||
|
||||
// for MTMD_GEN_PROCESS_TYPE_GEN_CODE
|
||||
const int32_t * codes;
|
||||
size_t n_codes;
|
||||
const float * embd; // the generated hidden state, to be fed back to backbone
|
||||
// it must have n_text_embd elements
|
||||
|
||||
// for MTMD_GEN_PROCESS_TYPE_GEN_WAV
|
||||
const float * audio;
|
||||
size_t n_samples;
|
||||
const char * state_data;
|
||||
size_t state_size;
|
||||
};
|
||||
// note: this API is stateless, caller must handle state management and audio frame accumulation
|
||||
MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx,
|
||||
const struct mtmd_gen_inp * inp,
|
||||
struct mtmd_gen_out * out);
|
||||
|
||||
/////////////////////////////////////////
|
||||
|
||||
// test function, to be used in test-mtmd-c-api.c
|
||||
|
||||
+193
-65
@@ -10,8 +10,10 @@
|
||||
#include <ctime>
|
||||
#include <atomic>
|
||||
#include <cstring>
|
||||
#include <cstdlib>
|
||||
#include <algorithm>
|
||||
#include <unordered_set>
|
||||
#include <tuple>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
|
||||
@@ -34,7 +36,40 @@ json server_tool::to_json() const {
|
||||
}
|
||||
|
||||
static constexpr size_t SERVER_TOOL_GIT_LS_FILES_MAX_OUTPUT = 8 * 1024 * 1024; // 8 MB
|
||||
static constexpr int SERVER_TOOL_GIT_LS_FILES_TIMEOUT = 15; // seconds
|
||||
// budget for one listing call, shared by the git and walker paths
|
||||
static constexpr int SERVER_TOOL_LIST_ENTRIES_TIMEOUT = 15; // seconds
|
||||
|
||||
// entry kinds a directory listing may return
|
||||
enum class list_kind {
|
||||
files, // regular files only
|
||||
dirs, // directories only
|
||||
all, // both
|
||||
};
|
||||
|
||||
// home directory, read once at first use (getenv is not thread safe against setenv)
|
||||
static const std::string & home_dir() {
|
||||
static const std::string home = [] {
|
||||
const char * h = getenv("HOME");
|
||||
#ifdef _WIN32
|
||||
if (h == nullptr) h = getenv("USERPROFILE");
|
||||
#endif
|
||||
return h ? std::string(h) : std::string();
|
||||
}();
|
||||
return home;
|
||||
}
|
||||
|
||||
static std::string expand_home(const std::string & path) {
|
||||
if (path.empty() || path[0] != '~') return path;
|
||||
if (path.size() > 1 && path[1] != '/' && path[1] != '\\') return path;
|
||||
const std::string & home = home_dir();
|
||||
if (home.empty()) return path;
|
||||
return home + path.substr(1);
|
||||
}
|
||||
|
||||
// depth of a '/'-separated relative path: "a/b/c" is 3
|
||||
static int entry_depth(const std::string & rel) {
|
||||
return 1 + (int) std::count(rel.begin(), rel.end(), '/');
|
||||
}
|
||||
|
||||
class tools_io {
|
||||
public:
|
||||
@@ -51,8 +86,17 @@ public:
|
||||
virtual bool file_size(const std::string & path, uintmax_t & out_size) const = 0;
|
||||
virtual bool read_file(const std::string & path, std::string & out) const = 0;
|
||||
virtual bool write_file(const std::string & path, const std::string & content) const = 0;
|
||||
// paths relative to `base`, '/'-separated; sets `err` if `base` isn't a directory
|
||||
virtual std::vector<std::string> list_files(const std::string & base, std::string & err) const = 0;
|
||||
// resolve `path` against the IO's working directory; absolute paths are returned unchanged
|
||||
virtual std::string resolve(const std::string & path) const = 0;
|
||||
struct list_entry {
|
||||
std::string rel; // '/'-separated, relative to `base`
|
||||
bool is_dir = false;
|
||||
};
|
||||
// entries relative to `base`; sets `err` if `base` isn't a directory
|
||||
// max_depth == 0 means unlimited, 1 means direct children of `base` only
|
||||
// `base` must already be resolved (absolute); `caller_path` is the path the
|
||||
// caller passed, used only for error messages
|
||||
virtual std::vector<list_entry> list_entries(const std::string & base, const std::string & caller_path, int max_depth, list_kind kind, std::string & err, bool & truncated) const = 0;
|
||||
// on_chunk, if set, is called with each chunk of output as it is read (before truncation cuts in);
|
||||
// returning false terminates the process early (e.g. the client disconnected)
|
||||
virtual exec_result run(
|
||||
@@ -67,6 +111,22 @@ public:
|
||||
// cwd, if non-empty, is used to resolve relative paths and as the working directory for run()
|
||||
explicit tools_io_basic(std::string cwd = "") : cwd(std::move(cwd)) {}
|
||||
|
||||
// expands a leading `~`, then resolves `path` against `cwd` (or the server
|
||||
// working directory when `cwd` is unset); the result is always absolute
|
||||
std::string resolve(const std::string & path) const override {
|
||||
std::string p = expand_home(path);
|
||||
if (fs::path(p).is_absolute()) {
|
||||
return p;
|
||||
}
|
||||
if (cwd.empty()) {
|
||||
std::error_code ec;
|
||||
fs::path cur = fs::current_path(ec);
|
||||
if (ec) return p;
|
||||
return (cur / p).string();
|
||||
}
|
||||
return (fs::path(cwd) / p).string();
|
||||
}
|
||||
|
||||
bool is_directory(const std::string & path) const override {
|
||||
std::error_code ec;
|
||||
return fs::is_directory(resolve(path), ec) && !ec;
|
||||
@@ -105,34 +165,41 @@ public:
|
||||
return (bool) f;
|
||||
}
|
||||
|
||||
std::vector<std::string> list_files(const std::string & base, std::string & err) const override {
|
||||
std::vector<list_entry> list_entries(const std::string & base, const std::string & caller_path, int max_depth, list_kind kind, std::string & err, bool & truncated) const override {
|
||||
err.clear();
|
||||
std::string abs_base = resolve(base);
|
||||
if (!is_directory(base)) {
|
||||
err = "path does not exist or is not a directory: " + base;
|
||||
truncated = false;
|
||||
std::error_code ec;
|
||||
if (!fs::is_directory(base, ec) || ec) {
|
||||
err = "path does not exist or is not a directory: " + caller_path;
|
||||
return {};
|
||||
}
|
||||
|
||||
auto res = run(
|
||||
{"git", "-C", abs_base, "ls-files", "--cached", "--others", "--exclude-standard"},
|
||||
SERVER_TOOL_GIT_LS_FILES_MAX_OUTPUT, SERVER_TOOL_GIT_LS_FILES_TIMEOUT);
|
||||
const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(SERVER_TOOL_LIST_ENTRIES_TIMEOUT);
|
||||
|
||||
if (res.exit_code == 0 && !res.timed_out) {
|
||||
std::vector<std::string> result;
|
||||
std::istringstream iss(res.output);
|
||||
std::string line;
|
||||
while (std::getline(iss, line)) {
|
||||
if (!line.empty() && line.back() == '\r') line.pop_back();
|
||||
if (line.empty()) continue;
|
||||
std::replace(line.begin(), line.end(), '\\', '/');
|
||||
if (is_regular_file((fs::path(base) / line).string())) {
|
||||
result.push_back(line);
|
||||
// git ls-files cannot list directories; use the walker when they are requested
|
||||
if (kind == list_kind::files) {
|
||||
auto res = run(
|
||||
{"git", "-C", base, "ls-files", "--cached", "--others", "--exclude-standard"},
|
||||
SERVER_TOOL_GIT_LS_FILES_MAX_OUTPUT, SERVER_TOOL_LIST_ENTRIES_TIMEOUT);
|
||||
|
||||
if (res.exit_code == 0 && !res.timed_out) {
|
||||
std::vector<list_entry> result;
|
||||
std::istringstream iss(res.output);
|
||||
std::string line;
|
||||
while (std::getline(iss, line)) {
|
||||
if (!line.empty() && line.back() == '\r') line.pop_back();
|
||||
if (line.empty()) continue;
|
||||
std::replace(line.begin(), line.end(), '\\', '/');
|
||||
if (max_depth > 0 && entry_depth(line) > max_depth) continue;
|
||||
if (is_regular_file((fs::path(base) / line).string())) {
|
||||
result.push_back({line, false});
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
return list_files_fallback(abs_base);
|
||||
return list_entries_fallback(base, max_depth, kind, deadline, truncated);
|
||||
}
|
||||
|
||||
exec_result run(
|
||||
@@ -211,14 +278,6 @@ public:
|
||||
private:
|
||||
std::string cwd;
|
||||
|
||||
// resolves `path` against `cwd` if `path` is relative and `cwd` is set; otherwise returns `path` unchanged
|
||||
std::string resolve(const std::string & path) const {
|
||||
if (cwd.empty() || fs::path(path).is_absolute()) {
|
||||
return path;
|
||||
}
|
||||
return (fs::path(cwd) / path).string();
|
||||
}
|
||||
|
||||
static const std::unordered_set<std::string> & junk_dir_names() {
|
||||
static const std::unordered_set<std::string> names = {
|
||||
".git", ".svn", ".hg", "node_modules", "__pycache__",
|
||||
@@ -227,28 +286,50 @@ private:
|
||||
return names;
|
||||
}
|
||||
|
||||
std::vector<std::string> list_files_fallback(const std::string & base) const {
|
||||
std::vector<std::string> result;
|
||||
std::vector<list_entry> list_entries_fallback(const std::string & base, int max_depth, list_kind kind,
|
||||
std::chrono::steady_clock::time_point deadline, bool & truncated) const {
|
||||
std::vector<list_entry> result;
|
||||
std::error_code ec;
|
||||
|
||||
std::vector<std::pair<fs::path, fs::path>> stack;
|
||||
stack.emplace_back(fs::path(base), fs::path());
|
||||
std::vector<std::tuple<fs::path, fs::path, int>> stack;
|
||||
stack.emplace_back(fs::path(base), fs::path(), 0);
|
||||
|
||||
while (!stack.empty()) {
|
||||
auto [dir, rel_dir] = stack.back();
|
||||
auto [dir, rel_dir, depth] = stack.back();
|
||||
stack.pop_back();
|
||||
|
||||
for (const auto & entry : fs::directory_iterator(dir, fs::directory_options::skip_permission_denied, ec)) {
|
||||
// the throwing increment would escape the tool on a directory that
|
||||
// goes away mid walk, so step the iterator explicitly
|
||||
fs::directory_iterator it(dir, fs::directory_options::skip_permission_denied, ec);
|
||||
for (const fs::directory_iterator end; it != end; it.increment(ec)) {
|
||||
if (ec) break;
|
||||
if (std::chrono::steady_clock::now() >= deadline) {
|
||||
truncated = true;
|
||||
return result;
|
||||
}
|
||||
const fs::directory_entry & entry = *it;
|
||||
std::string fname = entry.path().filename().string();
|
||||
std::error_code tec;
|
||||
if (entry.is_directory(tec)) {
|
||||
std::string rel = (rel_dir / fname).string();
|
||||
std::replace(rel.begin(), rel.end(), '\\', '/');
|
||||
if (kind == list_kind::dirs || kind == list_kind::all) {
|
||||
result.push_back({rel, true});
|
||||
}
|
||||
// junk directories stay selectable but are never walked: they
|
||||
// hold nothing worth searching and can be enormous
|
||||
if (junk_dir_names().count(fname) > 0) continue;
|
||||
stack.emplace_back(entry.path(), rel_dir / fname);
|
||||
// do not descend into symlinks: a link can point back to an
|
||||
// ancestor and loop forever
|
||||
if (!entry.is_symlink(tec) && (max_depth == 0 || depth + 1 < max_depth)) {
|
||||
stack.emplace_back(entry.path(), rel_dir / fname, depth + 1);
|
||||
}
|
||||
} else if (entry.is_regular_file(tec)) {
|
||||
std::string rel = (rel_dir / fname).string();
|
||||
std::replace(rel.begin(), rel.end(), '\\', '/');
|
||||
result.push_back(rel);
|
||||
if (kind == list_kind::files || kind == list_kind::all) {
|
||||
result.push_back({rel, false});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -363,6 +444,9 @@ struct server_tool_read_file : server_tool {
|
||||
//
|
||||
|
||||
static constexpr size_t SERVER_TOOL_FILE_SEARCH_MAX_RESULTS = 100;
|
||||
static constexpr const char * SERVER_TOOL_FILE_SEARCH_TYPE_FILE = "file";
|
||||
static constexpr const char * SERVER_TOOL_FILE_SEARCH_TYPE_DIR = "dir";
|
||||
static constexpr const char * SERVER_TOOL_FILE_SEARCH_TYPE_ALL = "all";
|
||||
|
||||
struct server_tool_file_glob_search : server_tool {
|
||||
server_tool_file_glob_search() {
|
||||
@@ -382,13 +466,18 @@ struct server_tool_file_glob_search : server_tool {
|
||||
"and common junk directories (.git, node_modules, build, dist, etc.) otherwise. "
|
||||
"A pattern with no '/' (e.g. \"*.cpp\") matches the file's basename at any depth. "
|
||||
"A pattern containing '/' matches the full relative path; unless already anchored with "
|
||||
"\"**/\" or a leading '/', it is automatically prefixed with \"**/\"."},
|
||||
"\"**/\" or a leading '/', it is automatically prefixed with \"**/\". "
|
||||
"Use type=\"dir\" or \"all\" to also list directories; directory entries are suffixed with '/' in the output. "
|
||||
"Note: directory listings do not apply .gitignore filtering."},
|
||||
{"parameters", {
|
||||
{"type", "object"},
|
||||
{"properties", {
|
||||
{"path", {{"type", "string"}, {"description", "Base directory to search in"}}},
|
||||
{"include", {{"type", "string"}, {"description", "Glob pattern for files to include (e.g. \"*.cpp\" or \"src/**/*.cpp\"). Default: **"}}},
|
||||
{"exclude", {{"type", "string"}, {"description", "Glob pattern for files to exclude"}}},
|
||||
{"path", {{"type", "string"}, {"description", "Base directory to search in"}}},
|
||||
{"include", {{"type", "string"}, {"description", "Glob pattern for files to include (e.g. \"*.cpp\" or \"src/**/*.cpp\"). Default: **"}}},
|
||||
{"exclude", {{"type", "string"}, {"description", "Glob pattern for files to exclude"}}},
|
||||
{"type", {{"type", "string"}, {"description", "Entry type to return: \"file\" (default), \"dir\" or \"all\""}}},
|
||||
{"max_depth", {{"type", "integer"}, {"description", "Maximum depth to descend into subdirectories (default: 0 = unlimited; 1 = direct children only)"}}},
|
||||
{"limit", {{"type", "integer"}, {"description", string_format("Maximum number of results to return (default %zu; values below 1 fall back to the default)", SERVER_TOOL_FILE_SEARCH_MAX_RESULTS)}}},
|
||||
}},
|
||||
{"required", json::array({"path"})},
|
||||
}},
|
||||
@@ -397,30 +486,56 @@ struct server_tool_file_glob_search : server_tool {
|
||||
}
|
||||
|
||||
json invoke(json params, server_tool::stream *) const override {
|
||||
std::string base = params.at("path").get<std::string>();
|
||||
std::string include = json_value(params, "include", std::string("**"));
|
||||
std::string exclude = json_value(params, "exclude", std::string(""));
|
||||
|
||||
auto io = make_tools_io(params);
|
||||
|
||||
std::string base = io->resolve(params.at("path").get<std::string>());
|
||||
// normalize to forward slashes so the web UI (which assumes '/') can
|
||||
// join the relative entries into absolute paths on Windows too
|
||||
std::replace(base.begin(), base.end(), '\\', '/');
|
||||
std::string include = json_value(params, "include", std::string("**"));
|
||||
std::string exclude = json_value(params, "exclude", std::string(""));
|
||||
std::string type = json_value(params, "type", std::string("file"));
|
||||
int max_depth = std::max(0, json_value(params, "max_depth", 0));
|
||||
int limit = json_value(params, "limit", (int) SERVER_TOOL_FILE_SEARCH_MAX_RESULTS);
|
||||
if (limit < 1) limit = SERVER_TOOL_FILE_SEARCH_MAX_RESULTS;
|
||||
limit = std::min(limit, (int) SERVER_TOOL_FILE_SEARCH_MAX_RESULTS);
|
||||
|
||||
list_kind kind;
|
||||
if (type == SERVER_TOOL_FILE_SEARCH_TYPE_FILE) {
|
||||
kind = list_kind::files;
|
||||
} else if (type == SERVER_TOOL_FILE_SEARCH_TYPE_DIR) {
|
||||
kind = list_kind::dirs;
|
||||
} else if (type == SERVER_TOOL_FILE_SEARCH_TYPE_ALL) {
|
||||
kind = list_kind::all;
|
||||
} else {
|
||||
return {{"error", "invalid type: " + type + " (expected \"file\", \"dir\" or \"all\")"}};
|
||||
}
|
||||
|
||||
std::string err;
|
||||
auto files = io->list_files(base, err);
|
||||
bool truncated = false;
|
||||
auto entries = io->list_entries(base, params.at("path").get<std::string>(), max_depth, kind, err, truncated);
|
||||
if (!err.empty()) {
|
||||
return {{"error", err}};
|
||||
}
|
||||
|
||||
std::vector<std::string> matches;
|
||||
for (const auto & rel : files) {
|
||||
if (!path_glob_match(include, rel)) continue;
|
||||
if (!exclude.empty() && path_glob_match(exclude, rel)) continue;
|
||||
matches.push_back(rel);
|
||||
std::vector<tools_io::list_entry> matches;
|
||||
for (const auto & entry : entries) {
|
||||
if (!path_glob_match(include, entry.rel)) continue;
|
||||
if (!exclude.empty() && path_glob_match(exclude, entry.rel)) continue;
|
||||
matches.push_back(entry);
|
||||
}
|
||||
|
||||
size_t total = matches.size();
|
||||
size_t shown = std::min(total, SERVER_TOOL_FILE_SEARCH_MAX_RESULTS);
|
||||
size_t shown = std::min(total, (size_t) limit);
|
||||
|
||||
std::ostringstream output_text;
|
||||
json entries_json = json::array();
|
||||
for (size_t i = 0; i < shown; i++) {
|
||||
output_text << matches[i] << "\n";
|
||||
output_text << matches[i].rel << (matches[i].is_dir ? "/" : "") << "\n";
|
||||
entries_json.push_back({
|
||||
{"path", matches[i].rel},
|
||||
{"type", matches[i].is_dir ? "dir" : "file"},
|
||||
});
|
||||
}
|
||||
|
||||
output_text << "\n---\nTotal matches: " << total << "\n";
|
||||
@@ -429,8 +544,16 @@ struct server_tool_file_glob_search : server_tool {
|
||||
"[%zu results limit reached (%zu total matches). Refine the glob pattern to narrow the search.]\n",
|
||||
shown, total);
|
||||
}
|
||||
if (truncated) {
|
||||
output_text << "[search timed out, results truncated]\n";
|
||||
}
|
||||
|
||||
return {{"plain_text_response", output_text.str()}};
|
||||
// `base` is always absolute (resolve falls back to the server cwd), so
|
||||
// API clients (e.g. the web UI picker) can join the relative entries
|
||||
// into absolute paths. `plain_text_response` is what the model sees;
|
||||
// `entries` is the same data as structured JSON for the UI picker,
|
||||
// which reads `entries`/`base` instead of re-parsing the text.
|
||||
return {{"plain_text_response", output_text.str()}, {"entries", entries_json}, {"base", base}};
|
||||
}
|
||||
};
|
||||
|
||||
@@ -513,18 +636,20 @@ struct server_tool_grep_search : server_tool {
|
||||
// collect (absolute_path, display_path) pairs to search
|
||||
std::vector<std::pair<std::string, std::string>> files;
|
||||
|
||||
if (io->is_regular_file(path)) {
|
||||
files.emplace_back(path, path);
|
||||
} else if (io->is_directory(path)) {
|
||||
const std::string abs_path = io->resolve(path);
|
||||
if (io->is_regular_file(abs_path)) {
|
||||
files.emplace_back(abs_path, path);
|
||||
} else if (io->is_directory(abs_path)) {
|
||||
std::string err;
|
||||
auto candidates = io->list_files(path, err);
|
||||
bool truncated = false;
|
||||
auto candidates = io->list_entries(abs_path, path, 0, list_kind::files, err, truncated);
|
||||
if (!err.empty()) {
|
||||
return {{"error", err}};
|
||||
}
|
||||
for (const auto & rel : candidates) {
|
||||
if (!path_glob_match(include, rel)) continue;
|
||||
if (!exclude.empty() && path_glob_match(exclude, rel)) continue;
|
||||
files.emplace_back((fs::path(path) / rel).string(), rel);
|
||||
for (const auto & entry : candidates) {
|
||||
if (!path_glob_match(include, entry.rel)) continue;
|
||||
if (!exclude.empty() && path_glob_match(exclude, entry.rel)) continue;
|
||||
files.emplace_back((fs::path(abs_path) / entry.rel).string(), entry.rel);
|
||||
}
|
||||
} else {
|
||||
return {{"error", "path does not exist: " + path}};
|
||||
@@ -1094,6 +1219,9 @@ struct server_tool_get_datetime : server_tool {
|
||||
// get_info: returns runtime info (OS name/version and cwd)
|
||||
//
|
||||
|
||||
static constexpr size_t SERVER_TOOL_GET_INFO_MAX_OUTPUT = 4096;
|
||||
static constexpr int SERVER_TOOL_GET_INFO_TIMEOUT = 5; // seconds
|
||||
|
||||
struct server_tool_get_info : server_tool {
|
||||
server_tool_get_info() {
|
||||
name = "get_info";
|
||||
@@ -1119,9 +1247,9 @@ struct server_tool_get_info : server_tool {
|
||||
auto io = make_tools_io(params);
|
||||
|
||||
#ifdef _WIN32
|
||||
auto res = io->run({"cmd", "/c", "ver"}, 4096, 5);
|
||||
auto res = io->run({"cmd", "/c", "ver"}, SERVER_TOOL_GET_INFO_MAX_OUTPUT, SERVER_TOOL_GET_INFO_TIMEOUT);
|
||||
#else
|
||||
auto res = io->run({"uname", "-a"}, 4096, 5);
|
||||
auto res = io->run({"uname", "-a"}, SERVER_TOOL_GET_INFO_MAX_OUTPUT, SERVER_TOOL_GET_INFO_TIMEOUT);
|
||||
#endif
|
||||
// "ver" prints a blank line before the version, so the output is stripped on both ends;
|
||||
// a failed spawn or a timeout leaves a diagnostic in res.output, which is not an OS name
|
||||
|
||||
@@ -164,3 +164,101 @@ def test_tools_builtin_edit_file_rejects_overlapping_edits():
|
||||
finally:
|
||||
if os.path.exists(log_path):
|
||||
os.remove(log_path)
|
||||
|
||||
|
||||
def test_tools_builtin_file_glob_search_type_dir(tmp_path):
|
||||
global server
|
||||
server.start()
|
||||
|
||||
(tmp_path / "project-alpha" / "src").mkdir(parents=True)
|
||||
(tmp_path / "project-alpha" / "README.md").write_text("alpha")
|
||||
(tmp_path / "project-alpha" / "src" / "main.cpp").write_text("int main() {}")
|
||||
(tmp_path / "project-beta").mkdir()
|
||||
(tmp_path / "project-beta" / "notes.txt").write_text("beta")
|
||||
|
||||
res = call_tool("file_glob_search", {"path": str(tmp_path), "type": "dir"})
|
||||
text = res["plain_text_response"]
|
||||
assert "project-alpha/" in text
|
||||
assert "project-beta/" in text
|
||||
assert "project-alpha/src/" in text
|
||||
assert "README.md" not in text
|
||||
types = {e["path"]: e["type"] for e in res["entries"]}
|
||||
assert types["project-alpha"] == "dir"
|
||||
assert types["project-alpha/src"] == "dir"
|
||||
|
||||
res_all = call_tool("file_glob_search", {"path": str(tmp_path), "type": "all", "include": "*proj*"})
|
||||
paths = [e["path"] for e in res_all["entries"]]
|
||||
assert "project-alpha" in paths
|
||||
assert "project-beta" in paths
|
||||
|
||||
|
||||
def test_tools_builtin_file_glob_search_max_depth_and_limit(tmp_path):
|
||||
global server
|
||||
server.start()
|
||||
|
||||
(tmp_path / "a" / "b" / "c").mkdir(parents=True)
|
||||
(tmp_path / "top.txt").write_text("top")
|
||||
(tmp_path / "a" / "mid.txt").write_text("mid")
|
||||
(tmp_path / "a" / "b" / "deep.txt").write_text("deep")
|
||||
|
||||
res = call_tool("file_glob_search", {"path": str(tmp_path), "max_depth": 1})
|
||||
assert "top.txt" in res["plain_text_response"]
|
||||
assert "mid.txt" not in res["plain_text_response"]
|
||||
|
||||
res = call_tool("file_glob_search", {"path": str(tmp_path), "max_depth": 2})
|
||||
assert "mid.txt" in res["plain_text_response"]
|
||||
assert "deep.txt" not in res["plain_text_response"]
|
||||
|
||||
res = call_tool("file_glob_search", {"path": str(tmp_path), "limit": 1})
|
||||
assert len(res["entries"]) == 1
|
||||
assert "Total matches: 3" in res["plain_text_response"]
|
||||
|
||||
|
||||
def test_tools_builtin_file_glob_search_rejects_invalid_type(tmp_path):
|
||||
global server
|
||||
server.start()
|
||||
|
||||
err = call_tool_expect_error("file_glob_search", {"path": str(tmp_path), "type": "bogus"})
|
||||
assert "invalid type" in err
|
||||
|
||||
|
||||
def test_tools_builtin_cwd_header_overrides_model_param(tmp_path):
|
||||
global server
|
||||
server.start()
|
||||
|
||||
workdir = tmp_path / "workdir"
|
||||
workdir.mkdir()
|
||||
(workdir / "marker.txt").write_text("marker")
|
||||
|
||||
# a model-provided "cwd" in the params is overridden by the x-tool-cwd header
|
||||
res = call_tool("read_file", {"path": "marker.txt", "cwd": "/definitely/not/a/real/path"},
|
||||
headers={"x-tool-cwd": str(workdir)})
|
||||
assert "marker" in res["plain_text_response"]
|
||||
|
||||
|
||||
def test_tools_builtin_cwd_relative_paths(tmp_path):
|
||||
global server
|
||||
server.start()
|
||||
|
||||
workdir = tmp_path / "workdir"
|
||||
workdir.mkdir()
|
||||
(workdir / "rel.txt").write_text("relative-content")
|
||||
|
||||
headers = {"x-tool-cwd": str(workdir)}
|
||||
|
||||
# relative paths in file tools resolve against the header cwd
|
||||
res = call_tool("read_file", {"path": "rel.txt"}, headers=headers)
|
||||
assert "relative-content" in res["plain_text_response"]
|
||||
|
||||
res = call_tool("write_file", {"path": "sub/out.txt", "content": "written"}, headers=headers)
|
||||
assert (workdir / "sub" / "out.txt").read_text() == "written"
|
||||
|
||||
res = call_tool("file_glob_search", {"path": ".", "include": "*.txt"}, headers=headers)
|
||||
assert "rel.txt" in res["plain_text_response"]
|
||||
|
||||
# absolute paths are unaffected by the cwd
|
||||
other = tmp_path / "other"
|
||||
other.mkdir()
|
||||
(other / "abs.txt").write_text("absolute-content")
|
||||
res = call_tool("read_file", {"path": str(other / "abs.txt")}, headers=headers)
|
||||
assert "absolute-content" in res["plain_text_response"]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
set(TARGET llama-tts)
|
||||
add_executable(${TARGET} tts.cpp)
|
||||
target_link_libraries(${TARGET} PRIVATE llama llama-common ${CMAKE_THREAD_LIBS_INIT})
|
||||
target_link_libraries(${TARGET} PRIVATE llama llama-common mtmd ${CMAKE_THREAD_LIBS_INIT})
|
||||
target_compile_features(${TARGET} PRIVATE cxx_std_17)
|
||||
|
||||
if(LLAMA_TOOLS_INSTALL)
|
||||
|
||||
+23
-106
@@ -1,117 +1,34 @@
|
||||
# llama.cpp/example/tts
|
||||
This example demonstrates the Text To Speech feature. It uses a
|
||||
[model](https://www.outeai.com/blog/outetts-0.2-500m) from
|
||||
[outeai](https://www.outeai.com/).
|
||||
# llama.cpp TTS
|
||||
|
||||
## Quickstart
|
||||
If you have built llama.cpp with SSL support you can simply run the
|
||||
following command and the required models will be downloaded automatically:
|
||||
```console
|
||||
$ build/bin/llama-tts --tts-oute-default -p "Hello world" && aplay output.wav
|
||||
```
|
||||
For details about the models and how to convert them to the required format
|
||||
see the following sections.
|
||||
This is a tool to demonstrate audio generation capability in llama.cpp via `libmtmd`. It was added via PR [#26254](https://github.com/ggml-org/llama.cpp/pull/26254)
|
||||
|
||||
### Model conversion
|
||||
Checkout or download the model that contains the LLM model:
|
||||
```console
|
||||
$ pushd models
|
||||
$ git clone --branch main --single-branch --depth 1 https://huggingface.co/OuteAI/OuteTTS-0.2-500M
|
||||
$ cd OuteTTS-0.2-500M && git lfs install && git lfs pull
|
||||
$ popd
|
||||
```
|
||||
Convert the model to .gguf format:
|
||||
```console
|
||||
(venv) python convert_hf_to_gguf.py models/OuteTTS-0.2-500M \
|
||||
--outfile models/outetts-0.2-0.5B-f16.gguf --outtype f16
|
||||
```
|
||||
The generated model will be `models/outetts-0.2-0.5B-f16.gguf`.
|
||||
Note: this tool used to serve as a demo for OuteTTS, but it was converted to a more model-agnostic tool.
|
||||
|
||||
We can optionally quantize this to Q8_0 using the following command:
|
||||
```console
|
||||
$ build/bin/llama-quantize models/outetts-0.2-0.5B-f16.gguf \
|
||||
models/outetts-0.2-0.5B-q8_0.gguf q8_0
|
||||
```
|
||||
The quantized model will be `models/outetts-0.2-0.5B-q8_0.gguf`.
|
||||
## Common usage
|
||||
|
||||
Next we do something similar for the audio decoder. First download or checkout
|
||||
the model for the voice decoder:
|
||||
```console
|
||||
$ pushd models
|
||||
$ git clone --branch main --single-branch --depth 1 https://huggingface.co/novateur/WavTokenizer-large-speech-75token
|
||||
$ cd WavTokenizer-large-speech-75token && git lfs install && git lfs pull
|
||||
$ popd
|
||||
```
|
||||
This model file is a PyTorch checkpoint (.ckpt) and we first need to convert it to
|
||||
huggingface format:
|
||||
```console
|
||||
(venv) python tools/tts/convert_pt_to_hf.py \
|
||||
models/WavTokenizer-large-speech-75token/wavtokenizer_large_speech_320_24k.ckpt
|
||||
...
|
||||
Model has been successfully converted and saved to models/WavTokenizer-large-speech-75token/model.safetensors
|
||||
Metadata has been saved to models/WavTokenizer-large-speech-75token/index.json
|
||||
Config has been saved to models/WavTokenizer-large-speech-75tokenconfig.json
|
||||
```
|
||||
Then we can convert the huggingface format to gguf:
|
||||
```console
|
||||
(venv) python convert_hf_to_gguf.py models/WavTokenizer-large-speech-75token \
|
||||
--outfile models/wavtokenizer-large-75-f16.gguf --outtype f16
|
||||
...
|
||||
INFO:hf-to-gguf:Model successfully exported to models/wavtokenizer-large-75-f16.gguf
|
||||
Simple usage:
|
||||
|
||||
```sh
|
||||
llama-tts -hf ggml-org/Qwen3-TTS-12Hz-1.7B-Base-GGUF -p "Hello world" --output out.wav
|
||||
```
|
||||
|
||||
### Running the example
|
||||
Common params:
|
||||
- Sampling params such as `--top-k`, `--top-p`, `--temp`, etc.
|
||||
- `-n <number_of_frames>` limits the output length, e.g. `-n 500`. Note that how many milliseconds each frame represents varies by model
|
||||
- Core inference params such as `-ngl`, `-b`, `-ub`, etc.
|
||||
|
||||
With both of the models generated, the LLM model and the voice decoder model,
|
||||
we can run the example:
|
||||
```console
|
||||
$ build/bin/llama-tts -m ./models/outetts-0.2-0.5B-q8_0.gguf \
|
||||
-mv ./models/wavtokenizer-large-75-f16.gguf \
|
||||
-p "Hello world"
|
||||
...
|
||||
main: audio written to file 'output.wav'
|
||||
```
|
||||
The output.wav file will contain the audio of the prompt. This can be heard
|
||||
by playing the file with a media player. On Linux the following command will
|
||||
play the audio:
|
||||
```console
|
||||
$ aplay output.wav
|
||||
```
|
||||
## Qwen3-TTS
|
||||
|
||||
### Running the example with llama-server
|
||||
Running this example with `llama-server` is also possible and requires two
|
||||
server instances to be started. One will serve the LLM model and the other
|
||||
will serve the voice decoder model.
|
||||
Available params:
|
||||
- `--tts-lang` can be `zh`, `en`, `de`, `it`, `pt`, `es`, `ja`, `ko`, `fr`, `ru` (default: `en`)
|
||||
- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3)
|
||||
|
||||
The LLM model server can be started with the following command:
|
||||
```console
|
||||
$ ./build/bin/llama-server -m ./models/outetts-0.2-0.5B-q8_0.gguf --port 8020
|
||||
```
|
||||
Example usage:
|
||||
|
||||
And the voice decoder model server can be started using:
|
||||
```console
|
||||
./build/bin/llama-server -m ./models/wavtokenizer-large-75-f16.gguf --port 8021 --embeddings --pooling none
|
||||
```
|
||||
|
||||
Then we can run [tts-outetts.py](tts-outetts.py) to generate the audio.
|
||||
|
||||
First create a virtual environment for python and install the required
|
||||
dependencies (this in only required to be done once):
|
||||
```console
|
||||
$ python3 -m venv venv
|
||||
$ source venv/bin/activate
|
||||
(venv) pip install requests numpy
|
||||
```
|
||||
|
||||
And then run the python script using:
|
||||
```conole
|
||||
(venv) python ./tools/tts/tts-outetts.py http://localhost:8020 http://localhost:8021 "Hello world"
|
||||
spectrogram generated: n_codes: 90, n_embd: 1282
|
||||
converting to audio ...
|
||||
audio generated: 28800 samples
|
||||
audio written to file "output.wav"
|
||||
```
|
||||
And to play the audio we can again use aplay or any other media player:
|
||||
```console
|
||||
$ aplay output.wav
|
||||
```sh
|
||||
llama-tts -hf ggml-org/Qwen3-TTS-12Hz-1.7B-Base-GGUF \
|
||||
-p "Hello world" \
|
||||
--tts-lang english \
|
||||
--tts-speaker-file speaker.mp3 \
|
||||
--output out.wav
|
||||
```
|
||||
|
||||
@@ -1,180 +0,0 @@
|
||||
# convert the https://huggingface.co/novateur/WavTokenizer-large-speech-75token to HF format
|
||||
# the goal is to be able to reuse the convert_hf_to_gguf.py after that to create a GGUF file with the WavTokenizer decoder
|
||||
#
|
||||
# TODO: this script is LLM-generated and probably very inefficient and should be rewritten
|
||||
|
||||
import torch
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import re
|
||||
|
||||
from safetensors.torch import save_file
|
||||
|
||||
# default
|
||||
model_path = './model.pt'
|
||||
|
||||
# read from CLI
|
||||
if len(sys.argv) > 1:
|
||||
model_path = sys.argv[1]
|
||||
|
||||
# get the directory of the input model
|
||||
path_dst = os.path.dirname(model_path)
|
||||
|
||||
print(f"Loading model from {model_path}")
|
||||
|
||||
model = torch.load(model_path, map_location='cpu')
|
||||
|
||||
#print(model)
|
||||
|
||||
# print all keys
|
||||
for key in model.keys():
|
||||
print(key)
|
||||
if key == 'hyper_parameters':
|
||||
#print(model[key])
|
||||
# dump as json pretty
|
||||
print(json.dumps(model[key], indent=4))
|
||||
#if key != 'state_dict' and key != 'optimizer_states':
|
||||
# print(model[key])
|
||||
|
||||
# Check if the loaded model is a state_dict or a model instance
|
||||
if isinstance(model, torch.nn.Module):
|
||||
state_dict = model.state_dict()
|
||||
else:
|
||||
state_dict = model
|
||||
|
||||
# Print the structure of the state_dict to understand its format
|
||||
print("State dictionary keys:")
|
||||
for key in state_dict.keys():
|
||||
print(key)
|
||||
|
||||
# Ensure the state_dict is flat and contains only torch.Tensor objects
|
||||
def flatten_state_dict(state_dict, parent_key='', sep='.'):
|
||||
items = []
|
||||
items_new = []
|
||||
|
||||
for k, v in state_dict.items():
|
||||
new_key = f"{parent_key}{sep}{k}" if parent_key else k
|
||||
if isinstance(v, torch.Tensor):
|
||||
items.append((new_key, v))
|
||||
elif isinstance(v, dict):
|
||||
items.extend(flatten_state_dict(v, new_key, sep=sep).items())
|
||||
return dict(items)
|
||||
|
||||
size_total_mb = 0
|
||||
|
||||
for key, value in list(items):
|
||||
# keep only what we need for inference
|
||||
if not key.startswith('state_dict.feature_extractor.encodec.quantizer.') and \
|
||||
not key.startswith('state_dict.backbone.') and \
|
||||
not key.startswith('state_dict.head.out'):
|
||||
print('Skipping key: ', key)
|
||||
continue
|
||||
|
||||
new_key = key
|
||||
|
||||
new_key = new_key.replace('state_dict.', '')
|
||||
new_key = new_key.replace('pos_net', 'posnet')
|
||||
|
||||
# check if matches "backbone.posnet.%d.bias" or "backbone.posnet.%d.weight"
|
||||
if new_key.startswith("backbone.posnet."):
|
||||
match = re.match(r"backbone\.posnet\.(\d+)\.(bias|weight)", new_key)
|
||||
if match:
|
||||
new_key = f"backbone.posnet.{match.group(1)}.norm.{match.group(2)}"
|
||||
|
||||
# "feature_extractor.encodec.quantizer.vq.layers.0._codebook.embed" -> "backbone.embedding.weight"
|
||||
if new_key == "feature_extractor.encodec.quantizer.vq.layers.0._codebook.embed":
|
||||
new_key = "backbone.embedding.weight"
|
||||
|
||||
# these are the only rows used
|
||||
# ref: https://github.com/edwko/OuteTTS/blob/a613e79c489d8256dd657ea9168d78de75895d82/outetts/wav_tokenizer/audio_codec.py#L100
|
||||
if new_key.endswith("norm.scale.weight"):
|
||||
new_key = new_key.replace("norm.scale.weight", "norm.weight")
|
||||
value = value[0]
|
||||
|
||||
if new_key.endswith("norm.shift.weight"):
|
||||
new_key = new_key.replace("norm.shift.weight", "norm.bias")
|
||||
value = value[0]
|
||||
|
||||
if new_key.endswith("gamma"):
|
||||
new_key = new_key.replace("gamma", "gamma.weight")
|
||||
|
||||
# convert from 1D [768] to 2D [768, 1] so that ggml_add can broadcast the bias
|
||||
if (new_key.endswith("norm.weight") or new_key.endswith("norm1.weight") or new_key.endswith("norm2.weight") or new_key.endswith(".bias")) and (new_key.startswith("backbone.posnet") or new_key.startswith("backbone.embed.bias")):
|
||||
value = value.unsqueeze(1)
|
||||
|
||||
if new_key.endswith("dwconv.bias"):
|
||||
value = value.unsqueeze(1)
|
||||
|
||||
size_mb = value.element_size() * value.nelement() / (1024 * 1024)
|
||||
print(f"{size_mb:8.2f} MB - {new_key}: {value.shape}")
|
||||
|
||||
size_total_mb += size_mb
|
||||
|
||||
#print(key, '->', new_key, ': ', value)
|
||||
#print(key, '->', new_key)
|
||||
|
||||
items_new.append((new_key, value))
|
||||
|
||||
print(f"Total size: {size_total_mb:8.2f} MB")
|
||||
|
||||
return dict(items_new)
|
||||
|
||||
flattened_state_dict = flatten_state_dict(state_dict)
|
||||
|
||||
|
||||
# Convert the model to the safetensors format
|
||||
output_path = path_dst + '/model.safetensors'
|
||||
save_file(flattened_state_dict, output_path)
|
||||
|
||||
print(f"Model has been successfully converted and saved to {output_path}")
|
||||
|
||||
# Calculate the total size of the .safetensors file
|
||||
total_size = os.path.getsize(output_path)
|
||||
|
||||
# Create the weight map
|
||||
weight_map = {
|
||||
"model.safetensors": ["*"] # Assuming all weights are in one file
|
||||
}
|
||||
|
||||
# Create metadata for the index.json file
|
||||
metadata = {
|
||||
"total_size": total_size,
|
||||
"weight_map": weight_map
|
||||
}
|
||||
|
||||
# Save the metadata to index.json
|
||||
index_path = path_dst + '/index.json'
|
||||
with open(index_path, 'w') as f:
|
||||
json.dump(metadata, f, indent=4)
|
||||
|
||||
print(f"Metadata has been saved to {index_path}")
|
||||
|
||||
config = {
|
||||
"architectures": [
|
||||
"WavTokenizerDec"
|
||||
],
|
||||
"hidden_size": 1282,
|
||||
"n_embd_features": 512,
|
||||
"n_ff": 2304,
|
||||
"vocab_size": 4096,
|
||||
"n_head": 1,
|
||||
"layer_norm_epsilon": 1e-6,
|
||||
"group_norm_epsilon": 1e-6,
|
||||
"group_norm_groups": 32,
|
||||
"max_position_embeddings": 8192, # ?
|
||||
"n_layer": 12,
|
||||
"posnet": {
|
||||
"n_embd": 768,
|
||||
"n_layer": 6
|
||||
},
|
||||
"convnext": {
|
||||
"n_embd": 768,
|
||||
"n_layer": 12
|
||||
},
|
||||
}
|
||||
|
||||
with open(path_dst + '/config.json', 'w') as f:
|
||||
json.dump(config, f, indent=4)
|
||||
|
||||
print(f"Config has been saved to {path_dst + 'config.json'}")
|
||||
@@ -1,299 +0,0 @@
|
||||
import sys
|
||||
#import json
|
||||
#import struct
|
||||
import requests
|
||||
import re
|
||||
import struct
|
||||
import numpy as np
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
|
||||
def fill_hann_window(size, periodic=True):
|
||||
if periodic:
|
||||
return np.hanning(size + 1)[:-1]
|
||||
return np.hanning(size)
|
||||
|
||||
|
||||
def irfft(n_fft, complex_input):
|
||||
return np.fft.irfft(complex_input, n=n_fft)
|
||||
|
||||
|
||||
def fold(buffer, n_out, n_win, n_hop, n_pad):
|
||||
result = np.zeros(n_out)
|
||||
n_frames = len(buffer) // n_win
|
||||
|
||||
for i in range(n_frames):
|
||||
start = i * n_hop
|
||||
end = start + n_win
|
||||
result[start:end] += buffer[i * n_win:(i + 1) * n_win]
|
||||
|
||||
return result[n_pad:-n_pad] if n_pad > 0 else result
|
||||
|
||||
|
||||
def process_frame(args):
|
||||
l, n_fft, ST, hann = args
|
||||
frame = irfft(n_fft, ST[l])
|
||||
frame = frame * hann
|
||||
hann2 = hann * hann
|
||||
return frame, hann2
|
||||
|
||||
|
||||
def embd_to_audio(embd, n_codes, n_embd, n_thread=4):
|
||||
embd = np.asarray(embd, dtype=np.float32).reshape(n_codes, n_embd)
|
||||
|
||||
n_fft = 1280
|
||||
n_hop = 320
|
||||
n_win = 1280
|
||||
n_pad = (n_win - n_hop) // 2
|
||||
n_out = (n_codes - 1) * n_hop + n_win
|
||||
|
||||
hann = fill_hann_window(n_fft, True)
|
||||
|
||||
E = np.zeros((n_embd, n_codes), dtype=np.float32)
|
||||
for l in range(n_codes):
|
||||
for k in range(n_embd):
|
||||
E[k, l] = embd[l, k]
|
||||
|
||||
half_embd = n_embd // 2
|
||||
S = np.zeros((n_codes, half_embd + 1), dtype=np.complex64)
|
||||
|
||||
for k in range(half_embd):
|
||||
for l in range(n_codes):
|
||||
mag = E[k, l]
|
||||
phi = E[k + half_embd, l]
|
||||
|
||||
mag = np.clip(np.exp(mag), 0, 1e2)
|
||||
S[l, k] = mag * np.exp(1j * phi)
|
||||
|
||||
res = np.zeros(n_codes * n_fft)
|
||||
hann2_buffer = np.zeros(n_codes * n_fft)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=n_thread) as executor:
|
||||
args = [(l, n_fft, S, hann) for l in range(n_codes)]
|
||||
results = list(executor.map(process_frame, args))
|
||||
|
||||
for l, (frame, hann2) in enumerate(results):
|
||||
res[l*n_fft:(l+1)*n_fft] = frame
|
||||
hann2_buffer[l*n_fft:(l+1)*n_fft] = hann2
|
||||
|
||||
audio = fold(res, n_out, n_win, n_hop, n_pad)
|
||||
env = fold(hann2_buffer, n_out, n_win, n_hop, n_pad)
|
||||
|
||||
mask = env > 1e-10
|
||||
audio[mask] /= env[mask]
|
||||
|
||||
return audio
|
||||
|
||||
|
||||
def save_wav(filename, audio_data, sample_rate):
|
||||
num_channels = 1
|
||||
bits_per_sample = 16
|
||||
bytes_per_sample = bits_per_sample // 8
|
||||
data_size = len(audio_data) * bytes_per_sample
|
||||
byte_rate = sample_rate * num_channels * bytes_per_sample
|
||||
block_align = num_channels * bytes_per_sample
|
||||
chunk_size = 36 + data_size # 36 = size of header minus first 8 bytes
|
||||
|
||||
header = struct.pack(
|
||||
'<4sI4s4sIHHIIHH4sI',
|
||||
b'RIFF',
|
||||
chunk_size,
|
||||
b'WAVE',
|
||||
b'fmt ',
|
||||
16, # fmt chunk size
|
||||
1, # audio format (PCM)
|
||||
num_channels,
|
||||
sample_rate,
|
||||
byte_rate,
|
||||
block_align,
|
||||
bits_per_sample,
|
||||
b'data',
|
||||
data_size
|
||||
)
|
||||
|
||||
audio_data = np.clip(audio_data * 32767, -32768, 32767)
|
||||
pcm_data = audio_data.astype(np.int16)
|
||||
|
||||
with open(filename, 'wb') as f:
|
||||
f.write(header)
|
||||
f.write(pcm_data.tobytes())
|
||||
|
||||
|
||||
def process_text(text: str):
|
||||
text = re.sub(r'\d+(\.\d+)?', lambda x: x.group(), text.lower()) # TODO this needs to be fixed
|
||||
text = re.sub(r'[-_/,\.\\]', ' ', text)
|
||||
text = re.sub(r'[^a-z\s]', '', text)
|
||||
text = re.sub(r'\s+', ' ', text).strip()
|
||||
return text.split()
|
||||
|
||||
# usage:
|
||||
# python tts-outetts.py http://server-llm:port http://server-dec:port "text"
|
||||
|
||||
if len(sys.argv) <= 3:
|
||||
print("usage: python tts-outetts.py http://server-llm:port http://server-dec:port \"text\"")
|
||||
exit(1)
|
||||
|
||||
host_llm = sys.argv[1]
|
||||
host_dec = sys.argv[2]
|
||||
text = sys.argv[3]
|
||||
|
||||
prefix = """<|im_start|>
|
||||
<|text_start|>the<|text_sep|>overall<|text_sep|>package<|text_sep|>from<|text_sep|>just<|text_sep|>two<|text_sep|>people<|text_sep|>is<|text_sep|>pretty<|text_sep|>remarkable<|text_sep|>sure<|text_sep|>i<|text_sep|>have<|text_sep|>some<|text_sep|>critiques<|text_sep|>about<|text_sep|>some<|text_sep|>of<|text_sep|>the<|text_sep|>gameplay<|text_sep|>aspects<|text_sep|>but<|text_sep|>its<|text_sep|>still<|text_sep|>really<|text_sep|>enjoyable<|text_sep|>and<|text_sep|>it<|text_sep|>looks<|text_sep|>lovely<|text_sep|>"""
|
||||
|
||||
words = process_text(text)
|
||||
words = "<|text_sep|>".join([i.strip() for i in words])
|
||||
words += "<|text_end|>\n"
|
||||
|
||||
# voice data
|
||||
# TODO: load from json
|
||||
#suffix = """<|audio_start|>
|
||||
#the<|t_0.08|><|code_start|><|257|><|740|><|636|><|913|><|788|><|1703|><|code_end|>
|
||||
#overall<|t_0.36|><|code_start|><|127|><|201|><|191|><|774|><|700|><|532|><|1056|><|557|><|798|><|298|><|1741|><|747|><|1662|><|1617|><|1702|><|1527|><|368|><|1588|><|1049|><|1008|><|1625|><|747|><|1576|><|728|><|1019|><|1696|><|1765|><|code_end|>
|
||||
#package<|t_0.56|><|code_start|><|935|><|584|><|1319|><|627|><|1016|><|1491|><|1344|><|1117|><|1526|><|1040|><|239|><|1435|><|951|><|498|><|723|><|1180|><|535|><|789|><|1649|><|1637|><|78|><|465|><|1668|><|901|><|595|><|1675|><|117|><|1009|><|1667|><|320|><|840|><|79|><|507|><|1762|><|1508|><|1228|><|1768|><|802|><|1450|><|1457|><|232|><|639|><|code_end|>
|
||||
#from<|t_0.19|><|code_start|><|604|><|782|><|1682|><|872|><|1532|><|1600|><|1036|><|1761|><|647|><|1554|><|1371|><|653|><|1595|><|950|><|code_end|>
|
||||
#just<|t_0.25|><|code_start|><|1782|><|1670|><|317|><|786|><|1748|><|631|><|599|><|1155|><|1364|><|1524|><|36|><|1591|><|889|><|1535|><|541|><|440|><|1532|><|50|><|870|><|code_end|>
|
||||
#two<|t_0.24|><|code_start|><|1681|><|1510|><|673|><|799|><|805|><|1342|><|330|><|519|><|62|><|640|><|1138|><|565|><|1552|><|1497|><|1552|><|572|><|1715|><|1732|><|code_end|>
|
||||
#people<|t_0.39|><|code_start|><|593|><|274|><|136|><|740|><|691|><|633|><|1484|><|1061|><|1138|><|1485|><|344|><|428|><|397|><|1562|><|645|><|917|><|1035|><|1449|><|1669|><|487|><|442|><|1484|><|1329|><|1832|><|1704|><|600|><|761|><|653|><|269|><|code_end|>
|
||||
#is<|t_0.16|><|code_start|><|566|><|583|><|1755|><|646|><|1337|><|709|><|802|><|1008|><|485|><|1583|><|652|><|10|><|code_end|>
|
||||
#pretty<|t_0.32|><|code_start|><|1818|><|1747|><|692|><|733|><|1010|><|534|><|406|><|1697|><|1053|><|1521|><|1355|><|1274|><|816|><|1398|><|211|><|1218|><|817|><|1472|><|1703|><|686|><|13|><|822|><|445|><|1068|><|code_end|>
|
||||
#remarkable<|t_0.68|><|code_start|><|230|><|1048|><|1705|><|355|><|706|><|1149|><|1535|><|1787|><|1356|><|1396|><|835|><|1583|><|486|><|1249|><|286|><|937|><|1076|><|1150|><|614|><|42|><|1058|><|705|><|681|><|798|><|934|><|490|><|514|><|1399|><|572|><|1446|><|1703|><|1346|><|1040|><|1426|><|1304|><|664|><|171|><|1530|><|625|><|64|><|1708|><|1830|><|1030|><|443|><|1509|><|1063|><|1605|><|1785|><|721|><|1440|><|923|><|code_end|>
|
||||
#sure<|t_0.36|><|code_start|><|792|><|1780|><|923|><|1640|><|265|><|261|><|1525|><|567|><|1491|><|1250|><|1730|><|362|><|919|><|1766|><|543|><|1|><|333|><|113|><|970|><|252|><|1606|><|133|><|302|><|1810|><|1046|><|1190|><|1675|><|code_end|>
|
||||
#i<|t_0.08|><|code_start|><|123|><|439|><|1074|><|705|><|1799|><|637|><|code_end|>
|
||||
#have<|t_0.16|><|code_start|><|1509|><|599|><|518|><|1170|><|552|><|1029|><|1267|><|864|><|419|><|143|><|1061|><|0|><|code_end|>
|
||||
#some<|t_0.16|><|code_start|><|619|><|400|><|1270|><|62|><|1370|><|1832|><|917|><|1661|><|167|><|269|><|1366|><|1508|><|code_end|>
|
||||
#critiques<|t_0.60|><|code_start|><|559|><|584|><|1163|><|1129|><|1313|><|1728|><|721|><|1146|><|1093|><|577|><|928|><|27|><|630|><|1080|><|1346|><|1337|><|320|><|1382|><|1175|><|1682|><|1556|><|990|><|1683|><|860|><|1721|><|110|><|786|><|376|><|1085|><|756|><|1523|><|234|><|1334|><|1506|><|1578|><|659|><|612|><|1108|><|1466|><|1647|><|308|><|1470|><|746|><|556|><|1061|><|code_end|>
|
||||
#about<|t_0.29|><|code_start|><|26|><|1649|><|545|><|1367|><|1263|><|1728|><|450|><|859|><|1434|><|497|><|1220|><|1285|><|179|><|755|><|1154|><|779|><|179|><|1229|><|1213|><|922|><|1774|><|1408|><|code_end|>
|
||||
#some<|t_0.23|><|code_start|><|986|><|28|><|1649|><|778|><|858|><|1519|><|1|><|18|><|26|><|1042|><|1174|><|1309|><|1499|><|1712|><|1692|><|1516|><|1574|><|code_end|>
|
||||
#of<|t_0.07|><|code_start|><|197|><|716|><|1039|><|1662|><|64|><|code_end|>
|
||||
#the<|t_0.08|><|code_start|><|1811|><|1568|><|569|><|886|><|1025|><|1374|><|code_end|>
|
||||
#gameplay<|t_0.48|><|code_start|><|1269|><|1092|><|933|><|1362|><|1762|><|1700|><|1675|><|215|><|781|><|1086|><|461|><|838|><|1022|><|759|><|649|><|1416|><|1004|><|551|><|909|><|787|><|343|><|830|><|1391|><|1040|><|1622|><|1779|><|1360|><|1231|><|1187|><|1317|><|76|><|997|><|989|><|978|><|737|><|189|><|code_end|>
|
||||
#aspects<|t_0.56|><|code_start|><|1423|><|797|><|1316|><|1222|><|147|><|719|><|1347|><|386|><|1390|><|1558|><|154|><|440|><|634|><|592|><|1097|><|1718|><|712|><|763|><|1118|><|1721|><|1311|><|868|><|580|><|362|><|1435|><|868|><|247|><|221|><|886|><|1145|><|1274|><|1284|><|457|><|1043|><|1459|><|1818|><|62|><|599|><|1035|><|62|><|1649|><|778|><|code_end|>
|
||||
#but<|t_0.20|><|code_start|><|780|><|1825|><|1681|><|1007|><|861|><|710|><|702|><|939|><|1669|><|1491|><|613|><|1739|><|823|><|1469|><|648|><|code_end|>
|
||||
#its<|t_0.09|><|code_start|><|92|><|688|><|1623|><|962|><|1670|><|527|><|599|><|code_end|>
|
||||
#still<|t_0.27|><|code_start|><|636|><|10|><|1217|><|344|><|713|><|957|><|823|><|154|><|1649|><|1286|><|508|><|214|><|1760|><|1250|><|456|><|1352|><|1368|><|921|><|615|><|5|><|code_end|>
|
||||
#really<|t_0.36|><|code_start|><|55|><|420|><|1008|><|1659|><|27|><|644|><|1266|><|617|><|761|><|1712|><|109|><|1465|><|1587|><|503|><|1541|><|619|><|197|><|1019|><|817|><|269|><|377|><|362|><|1381|><|507|><|1488|><|4|><|1695|><|code_end|>
|
||||
#enjoyable<|t_0.49|><|code_start|><|678|><|501|><|864|><|319|><|288|><|1472|><|1341|><|686|><|562|><|1463|><|619|><|1563|><|471|><|911|><|730|><|1811|><|1006|><|520|><|861|><|1274|><|125|><|1431|><|638|><|621|><|153|><|876|><|1770|><|437|><|987|><|1653|><|1109|><|898|><|1285|><|80|><|593|><|1709|><|843|><|code_end|>
|
||||
#and<|t_0.15|><|code_start|><|1285|><|987|><|303|><|1037|><|730|><|1164|><|502|><|120|><|1737|><|1655|><|1318|><|code_end|>
|
||||
#it<|t_0.09|><|code_start|><|848|><|1366|><|395|><|1601|><|1513|><|593|><|1302|><|code_end|>
|
||||
#looks<|t_0.27|><|code_start|><|1281|><|1266|><|1755|><|572|><|248|><|1751|><|1257|><|695|><|1380|><|457|><|659|><|585|><|1315|><|1105|><|1776|><|736|><|24|><|736|><|654|><|1027|><|code_end|>
|
||||
#lovely<|t_0.56|><|code_start|><|634|><|596|><|1766|><|1556|><|1306|><|1285|><|1481|><|1721|><|1123|><|438|><|1246|><|1251|><|795|><|659|><|1381|><|1658|><|217|><|1772|><|562|><|952|><|107|><|1129|><|1112|><|467|><|550|><|1079|><|840|><|1615|><|1469|><|1380|><|168|><|917|><|836|><|1827|><|437|><|583|><|67|><|595|><|1087|><|1646|><|1493|><|1677|><|code_end|>"""
|
||||
|
||||
# TODO: tokenization is slow for some reason - here is pre-tokenized input
|
||||
suffix = [ 151667, 198, 1782, 155780, 151669, 151929, 152412, 152308, 152585, 152460, 153375, 151670, 198, 74455,
|
||||
155808, 151669, 151799, 151873, 151863, 152446, 152372, 152204, 152728, 152229, 152470, 151970, 153413,
|
||||
152419, 153334, 153289, 153374, 153199, 152040, 153260, 152721, 152680, 153297, 152419, 153248, 152400,
|
||||
152691, 153368, 153437, 151670, 198, 1722, 155828, 151669, 152607, 152256, 152991, 152299, 152688, 153163,
|
||||
153016, 152789, 153198, 152712, 151911, 153107, 152623, 152170, 152395, 152852, 152207, 152461, 153321,
|
||||
153309, 151750, 152137, 153340, 152573, 152267, 153347, 151789, 152681, 153339, 151992, 152512, 151751,
|
||||
152179, 153434, 153180, 152900, 153440, 152474, 153122, 153129, 151904, 152311, 151670, 198, 1499, 155791,
|
||||
151669, 152276, 152454, 153354, 152544, 153204, 153272, 152708, 153433, 152319, 153226, 153043, 152325,
|
||||
153267, 152622, 151670, 198, 4250, 155797, 151669, 153454, 153342, 151989, 152458, 153420, 152303, 152271,
|
||||
152827, 153036, 153196, 151708, 153263, 152561, 153207, 152213, 152112, 153204, 151722, 152542, 151670, 198,
|
||||
19789, 155796, 151669, 153353, 153182, 152345, 152471, 152477, 153014, 152002, 152191, 151734, 152312, 152810,
|
||||
152237, 153224, 153169, 153224, 152244, 153387, 153404, 151670, 198, 16069, 155811, 151669, 152265, 151946,
|
||||
151808, 152412, 152363, 152305, 153156, 152733, 152810, 153157, 152016, 152100, 152069, 153234, 152317,
|
||||
152589, 152707, 153121, 153341, 152159, 152114, 153156, 153001, 153504, 153376, 152272, 152433, 152325,
|
||||
151941, 151670, 198, 285, 155788, 151669, 152238, 152255, 153427, 152318, 153009, 152381, 152474, 152680,
|
||||
152157, 153255, 152324, 151682, 151670, 198, 32955, 155804, 151669, 153490, 153419, 152364, 152405, 152682,
|
||||
152206, 152078, 153369, 152725, 153193, 153027, 152946, 152488, 153070, 151883, 152890, 152489, 153144,
|
||||
153375, 152358, 151685, 152494, 152117, 152740, 151670, 198, 37448, 480, 155840, 151669, 151902, 152720,
|
||||
153377, 152027, 152378, 152821, 153207, 153459, 153028, 153068, 152507, 153255, 152158, 152921, 151958,
|
||||
152609, 152748, 152822, 152286, 151714, 152730, 152377, 152353, 152470, 152606, 152162, 152186, 153071,
|
||||
152244, 153118, 153375, 153018, 152712, 153098, 152976, 152336, 151843, 153202, 152297, 151736, 153380,
|
||||
153502, 152702, 152115, 153181, 152735, 153277, 153457, 152393, 153112, 152595, 151670, 198, 19098, 155808,
|
||||
151669, 152464, 153452, 152595, 153312, 151937, 151933, 153197, 152239, 153163, 152922, 153402, 152034,
|
||||
152591, 153438, 152215, 151673, 152005, 151785, 152642, 151924, 153278, 151805, 151974, 153482, 152718,
|
||||
152862, 153347, 151670, 198, 72, 155780, 151669, 151795, 152111, 152746, 152377, 153471, 152309, 151670, 198,
|
||||
19016, 155788, 151669, 153181, 152271, 152190, 152842, 152224, 152701, 152939, 152536, 152091, 151815, 152733,
|
||||
151672, 151670, 198, 14689, 155788, 151669, 152291, 152072, 152942, 151734, 153042, 153504, 152589, 153333,
|
||||
151839, 151941, 153038, 153180, 151670, 198, 36996, 8303, 155832, 151669, 152231, 152256, 152835, 152801,
|
||||
152985, 153400, 152393, 152818, 152765, 152249, 152600, 151699, 152302, 152752, 153018, 153009, 151992,
|
||||
153054, 152847, 153354, 153228, 152662, 153355, 152532, 153393, 151782, 152458, 152048, 152757, 152428,
|
||||
153195, 151906, 153006, 153178, 153250, 152331, 152284, 152780, 153138, 153319, 151980, 153142, 152418,
|
||||
152228, 152733, 151670, 198, 9096, 155801, 151669, 151698, 153321, 152217, 153039, 152935, 153400, 152122,
|
||||
152531, 153106, 152169, 152892, 152957, 151851, 152427, 152826, 152451, 151851, 152901, 152885, 152594,
|
||||
153446, 153080, 151670, 198, 14689, 155795, 151669, 152658, 151700, 153321, 152450, 152530, 153191, 151673,
|
||||
151690, 151698, 152714, 152846, 152981, 153171, 153384, 153364, 153188, 153246, 151670, 198, 1055, 155779,
|
||||
151669, 151869, 152388, 152711, 153334, 151736, 151670, 198, 1782, 155780, 151669, 153483, 153240, 152241,
|
||||
152558, 152697, 153046, 151670, 198, 5804, 1363, 155820, 151669, 152941, 152764, 152605, 153034, 153434,
|
||||
153372, 153347, 151887, 152453, 152758, 152133, 152510, 152694, 152431, 152321, 153088, 152676, 152223,
|
||||
152581, 152459, 152015, 152502, 153063, 152712, 153294, 153451, 153032, 152903, 152859, 152989, 151748,
|
||||
152669, 152661, 152650, 152409, 151861, 151670, 198, 300, 7973, 155828, 151669, 153095, 152469, 152988,
|
||||
152894, 151819, 152391, 153019, 152058, 153062, 153230, 151826, 152112, 152306, 152264, 152769, 153390,
|
||||
152384, 152435, 152790, 153393, 152983, 152540, 152252, 152034, 153107, 152540, 151919, 151893, 152558,
|
||||
152817, 152946, 152956, 152129, 152715, 153131, 153490, 151734, 152271, 152707, 151734, 153321, 152450,
|
||||
151670, 198, 8088, 155792, 151669, 152452, 153497, 153353, 152679, 152533, 152382, 152374, 152611, 153341,
|
||||
153163, 152285, 153411, 152495, 153141, 152320, 151670, 198, 1199, 155781, 151669, 151764, 152360, 153295,
|
||||
152634, 153342, 152199, 152271, 151670, 198, 43366, 155799, 151669, 152308, 151682, 152889, 152016, 152385,
|
||||
152629, 152495, 151826, 153321, 152958, 152180, 151886, 153432, 152922, 152128, 153024, 153040, 152593,
|
||||
152287, 151677, 151670, 198, 53660, 155808, 151669, 151727, 152092, 152680, 153331, 151699, 152316, 152938,
|
||||
152289, 152433, 153384, 151781, 153137, 153259, 152175, 153213, 152291, 151869, 152691, 152489, 151941,
|
||||
152049, 152034, 153053, 152179, 153160, 151676, 153367, 151670, 198, 268, 4123, 480, 155821, 151669, 152350,
|
||||
152173, 152536, 151991, 151960, 153144, 153013, 152358, 152234, 153135, 152291, 153235, 152143, 152583,
|
||||
152402, 153483, 152678, 152192, 152533, 152946, 151797, 153103, 152310, 152293, 151825, 152548, 153442,
|
||||
152109, 152659, 153325, 152781, 152570, 152957, 151752, 152265, 153381, 152515, 151670, 198, 437, 155787,
|
||||
151669, 152957, 152659, 151975, 152709, 152402, 152836, 152174, 151792, 153409, 153327, 152990, 151670, 198,
|
||||
275, 155781, 151669, 152520, 153038, 152067, 153273, 153185, 152265, 152974, 151670, 198, 94273, 155799,
|
||||
151669, 152953, 152938, 153427, 152244, 151920, 153423, 152929, 152367, 153052, 152129, 152331, 152257,
|
||||
152987, 152777, 153448, 152408, 151696, 152408, 152326, 152699, 151670, 198, 385, 16239, 155828, 151669,
|
||||
152306, 152268, 153438, 153228, 152978, 152957, 153153, 153393, 152795, 152110, 152918, 152923, 152467,
|
||||
152331, 153053, 153330, 151889, 153444, 152234, 152624, 151779, 152801, 152784, 152139, 152222, 152751,
|
||||
152512, 153287, 153141, 153052, 151840, 152589, 152508, 153499, 152109, 152255, 151739, 152267, 152759,
|
||||
153318, 153165, 153349, 151670, ]
|
||||
|
||||
response = requests.post(
|
||||
host_llm + "/completion",
|
||||
json={
|
||||
"prompt": [prefix + words, *suffix],
|
||||
"n_predict": 1024,
|
||||
"cache_prompt": True,
|
||||
"return_tokens": True,
|
||||
"samplers": ["top_k"],
|
||||
"top_k": 16,
|
||||
"seed": 1003,
|
||||
}
|
||||
)
|
||||
|
||||
response_json = response.json()
|
||||
|
||||
#print(json.dumps(response_json, indent=4))
|
||||
#print(json.dumps(response_json["prompt"], indent=4).replace("\\n", "\n"))
|
||||
#print(json.dumps(response_json["timings"], indent=4))
|
||||
#print(json.dumps(response_json["tokens"], indent=4))
|
||||
|
||||
codes = response_json["tokens"]
|
||||
|
||||
codes = [t - 151672 for t in codes if t >= 151672 and t <= 155772]
|
||||
|
||||
response = requests.post(
|
||||
host_dec + "/embeddings",
|
||||
json={
|
||||
"input": [*codes],
|
||||
}
|
||||
)
|
||||
|
||||
response_json = response.json()
|
||||
|
||||
#print(json.dumps(response_json, indent=4))
|
||||
|
||||
# spectrogram
|
||||
embd = response_json[0]["embedding"]
|
||||
|
||||
n_codes = len(embd)
|
||||
n_embd = len(embd[0])
|
||||
|
||||
print('spectrogram generated: n_codes: %d, n_embd: %d' % (n_codes, n_embd))
|
||||
|
||||
# post-process the spectrogram to convert to audio
|
||||
print('converting to audio ...')
|
||||
audio = embd_to_audio(embd, n_codes, n_embd)
|
||||
print('audio generated: %d samples' % len(audio))
|
||||
|
||||
filename = "output.wav"
|
||||
sample_rate = 24000 # sampling rate
|
||||
|
||||
# zero out first 0.25 seconds
|
||||
audio[:24000 // 4] = 0.0
|
||||
|
||||
save_wav(filename, audio, sample_rate)
|
||||
print('audio written to file "%s"' % filename)
|
||||
+152
-1043
File diff suppressed because it is too large
Load Diff
Vendored
+9
@@ -142,5 +142,14 @@ declare global {
|
||||
interface Window {
|
||||
idxThemeStyle?: number;
|
||||
idxCodeBlock?: number;
|
||||
|
||||
// File System Access API - missing from older DOM lib versions.
|
||||
// Used by ChatFormWorkingDirectory's native folder picker. Feature availability
|
||||
// is gated at runtime via `typeof window.showDirectoryPicker === 'function'`.
|
||||
showDirectoryPicker: (options?: {
|
||||
id?: string;
|
||||
mode?: 'read' | 'readwrite';
|
||||
startIn?: FileSystemHandle | string;
|
||||
}) => Promise<FileSystemDirectoryHandle>;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
ChatFormMcpResourcesList,
|
||||
ChatFormPickers,
|
||||
ChatFormTextarea,
|
||||
ChatFormWorkingDirectory,
|
||||
DialogMcpResourcesBrowser
|
||||
} from '$lib/components/app';
|
||||
import {
|
||||
@@ -31,7 +32,13 @@
|
||||
import { chatStore } from '$lib/stores/chat.svelte';
|
||||
import { mcpStore } from '$lib/stores/mcp.svelte';
|
||||
import { mcpHasResourceAttachments } from '$lib/stores/mcp-resources.svelte';
|
||||
import { conversationsStore, activeMessages } from '$lib/stores/conversations.svelte';
|
||||
import { toolsStore } from '$lib/stores/tools.svelte';
|
||||
import {
|
||||
conversationsStore,
|
||||
activeMessages,
|
||||
activeConversation,
|
||||
pendingCwd
|
||||
} from '$lib/stores/conversations.svelte';
|
||||
import type { GetPromptResult, MCPPromptInfo, MCPResourceInfo, PromptMessage } from '$lib/types';
|
||||
import { isIMEComposing, parseClipboardContent, uuid } from '$lib/utils';
|
||||
import {
|
||||
@@ -107,6 +114,15 @@
|
||||
let isInlineResourcePickerOpen = $state(false);
|
||||
let resourceSearchQuery = $state('');
|
||||
|
||||
let cwd = $derived(activeConversation()?.cwd ?? pendingCwd());
|
||||
|
||||
async function handleWorkingDirectoryChange(value: string | null) {
|
||||
await conversationsStore.setCwd(value);
|
||||
if (conversationsStore.activeConversation) {
|
||||
await chatStore.recordCwdChange(value?.trim() || null);
|
||||
}
|
||||
}
|
||||
|
||||
// Resource Dialog State
|
||||
let isResourceDialogOpen = $state(false);
|
||||
let preSelectedResourceUri = $state<string | undefined>(undefined);
|
||||
@@ -155,6 +171,12 @@
|
||||
audioRecorder = new AudioRecorder();
|
||||
});
|
||||
|
||||
// Defer so the closing popover's focus scope tears down first - bits-ui
|
||||
// yanks a synchronous focus() back into the still-mounted popover.
|
||||
function refocusInput() {
|
||||
queueMicrotask(() => textareaRef?.focus());
|
||||
}
|
||||
|
||||
export function focus() {
|
||||
textareaRef?.focus();
|
||||
}
|
||||
@@ -470,7 +492,7 @@
|
||||
<ChatFormFileInputInvisible bind:this={fileInputRef} onFileSelect={handleFileSelect} />
|
||||
|
||||
<form
|
||||
class="relative {className}"
|
||||
class="relative grid {className}"
|
||||
onsubmit={(event) => {
|
||||
event.preventDefault();
|
||||
|
||||
@@ -559,6 +581,15 @@
|
||||
</div>
|
||||
|
||||
<ContextGaugePopup />
|
||||
|
||||
{#if toolsStore.builtinTools.length > 0}
|
||||
<ChatFormWorkingDirectory
|
||||
directory={cwd}
|
||||
onChange={handleWorkingDirectoryChange}
|
||||
onClose={refocusInput}
|
||||
{disabled}
|
||||
/>
|
||||
{/if}
|
||||
</form>
|
||||
|
||||
<DialogMcpResourcesBrowser
|
||||
|
||||
@@ -0,0 +1,479 @@
|
||||
<script lang="ts">
|
||||
import { FolderOpen } from '@lucide/svelte';
|
||||
import { untrack } from 'svelte';
|
||||
import { SvelteMap } from 'svelte/reactivity';
|
||||
import { ToolsService } from '$lib/services/tools.service';
|
||||
import { toolsStore } from '$lib/stores/tools.svelte';
|
||||
import { BuiltInTool, GlobSearchType, KeyboardKey } from '$lib/enums';
|
||||
import {
|
||||
abbreviateHome,
|
||||
buildCaseInsensitiveGlob,
|
||||
joinPath,
|
||||
lastPathSegment,
|
||||
rankEntries,
|
||||
splitPathQuery,
|
||||
type GlobEntry
|
||||
} from '$lib/utils';
|
||||
import { debounce } from '$lib/utils/debounce';
|
||||
import * as Popover from '$lib/components/ui/popover';
|
||||
import SearchInput from '$lib/components/app/forms/SearchInput.svelte';
|
||||
import ChatFormWorkingDirectoryChip from './ChatFormWorkingDirectoryChip.svelte';
|
||||
import ChatFormWorkingDirectoryResultsList from './ChatFormWorkingDirectoryResultsList.svelte';
|
||||
import {
|
||||
DEFAULT_MOBILE_BREAKPOINT,
|
||||
GLOB_WILDCARD,
|
||||
HOME_TILDE,
|
||||
MAX_RESULTS_SHOWN,
|
||||
NATIVE_LIMIT,
|
||||
NATIVE_MAX_DEPTH,
|
||||
PATH_NAV_MAX_DEPTH,
|
||||
SEARCH_DEBOUNCE_MS,
|
||||
SEARCH_LIMIT,
|
||||
SEARCH_MAX_DEPTH
|
||||
} from '$lib/constants';
|
||||
|
||||
// Microtask delay so the popover's focus scope tears down first.
|
||||
const FOCUS_DELAY_MS = 0;
|
||||
|
||||
interface Props {
|
||||
class?: string;
|
||||
disabled?: boolean;
|
||||
directory?: string | null;
|
||||
onChange?: (directory: string | null) => void;
|
||||
/**
|
||||
* Lets the host refocus the chat input so typing can resume without
|
||||
* an extra click after the popover closes.
|
||||
*/
|
||||
onClose?: () => void;
|
||||
}
|
||||
|
||||
let {
|
||||
class: className = '',
|
||||
disabled = false,
|
||||
directory = $bindable(null),
|
||||
onChange,
|
||||
onClose
|
||||
}: Props = $props();
|
||||
|
||||
// File System Access API is opt-in: when available (Chrome / Edge / Opera) the popover
|
||||
// exposes a "Browse" button that opens the native folder picker. When unavailable the
|
||||
// popover still works via the text input - no alerts, no upload semantics.
|
||||
const pickerSupported =
|
||||
typeof window !== 'undefined' && typeof window.showDirectoryPicker === 'function';
|
||||
|
||||
// Popover open state; the element handles outside-click and Escape.
|
||||
let isOpen = $state(false);
|
||||
let inputValue = $state('');
|
||||
let searchInputRef: HTMLInputElement | null = $state(null);
|
||||
|
||||
let queryResults = $state<string[]>([]);
|
||||
let isSearching = $state(false);
|
||||
let searchError = $state<string | null>(null);
|
||||
let hoveredIndex = $state(-1);
|
||||
// Bumped only by ArrowUp/ArrowDown handlers; the list scrolls the
|
||||
// highlighted row into view only via this trigger, never on hover.
|
||||
let scrollTrigger = $state(0);
|
||||
let listContainer = $state<HTMLDivElement | null>(null);
|
||||
|
||||
// Absolute home directory on the server, resolved once per session by
|
||||
// the tools store. Anchors both the search scope and the chip's `~`
|
||||
// abbreviation.
|
||||
let homeBase = $derived(toolsStore.serverHome);
|
||||
|
||||
// AbortController + sequence counter to discard stale responses when the user
|
||||
// keeps typing; a newer call aborts the previous one. The sequence counter
|
||||
// also covers the gap between abort and the catch handler.
|
||||
let searchController: AbortController | null = null;
|
||||
let searchSeq = 0;
|
||||
|
||||
// Cache of the last file_glob_search result per (parent, include, max_depth),
|
||||
// so repeated queries in the same directory don't re-walk the tree. Entries
|
||||
// expire after a short TTL.
|
||||
const SEARCH_CACHE_TTL_MS = 2000;
|
||||
const searchCache = new SvelteMap<string, { results: GlobEntry[]; base: string; at: number }>();
|
||||
|
||||
const runSearch = debounce((query: string) => {
|
||||
void doSearch(query);
|
||||
}, SEARCH_DEBOUNCE_MS);
|
||||
|
||||
// Resolve home eagerly on mount so the chip can abbreviate before the
|
||||
// user opens the picker. resolveServerHome() is cached, so repeat calls
|
||||
// (e.g. from handleOpenChange) are no-ops.
|
||||
$effect(() => {
|
||||
if (typeof window === 'undefined') return;
|
||||
void toolsStore.resolveServerHome();
|
||||
});
|
||||
|
||||
// Auto-focus the search input when the popover opens.
|
||||
// HTML `autofocus` is unreliable on dynamically shown elements, so we
|
||||
// use a microtask (0ms setTimeout) after the effect flushes.
|
||||
$effect(() => {
|
||||
if (!isOpen) return;
|
||||
setTimeout(() => searchInputRef?.focus(), FOCUS_DELAY_MS);
|
||||
});
|
||||
|
||||
let lastScrollTrigger: number | null = null;
|
||||
|
||||
// hoveredIndex/queryResults are untracked so hover and result replacement
|
||||
// never re-fire the scroll; keyboard nav is the only path that bumps the trigger
|
||||
$effect(() => {
|
||||
if (scrollTrigger === lastScrollTrigger) return;
|
||||
lastScrollTrigger = scrollTrigger;
|
||||
untrack(() => {
|
||||
if (!listContainer) return;
|
||||
if (hoveredIndex < 0 || hoveredIndex >= queryResults.length) return;
|
||||
const selectedElement = listContainer.querySelector(
|
||||
`[data-result-index="${hoveredIndex}"]`
|
||||
) as HTMLElement | null;
|
||||
selectedElement?.scrollIntoView({ block: 'nearest', inline: 'nearest' });
|
||||
});
|
||||
});
|
||||
|
||||
function cancelSearch() {
|
||||
searchController?.abort();
|
||||
searchSeq++;
|
||||
isSearching = false;
|
||||
}
|
||||
|
||||
// Effective directory the current search runs against (shown in the
|
||||
// footer); updated by doSearch, including when an exactly-typed
|
||||
// directory is "entered".
|
||||
let searchScope = $state(HOME_TILDE);
|
||||
|
||||
// Runs a directory listing through the cache, so a repeated query in the
|
||||
// same directory does not re-walk the tree on the server.
|
||||
async function searchDirs(
|
||||
path: string,
|
||||
include: string,
|
||||
maxDepth: number,
|
||||
signal: AbortSignal
|
||||
): Promise<{ base: string; entries: GlobEntry[]; error?: string }> {
|
||||
const key = `${path}\u0000${include}\u0000${maxDepth}`;
|
||||
const cached = searchCache.get(key);
|
||||
if (cached && Date.now() - cached.at < SEARCH_CACHE_TTL_MS) {
|
||||
return { base: cached.base, entries: cached.results };
|
||||
}
|
||||
const res = await ToolsService.executeToolRaw(
|
||||
BuiltInTool.FILE_GLOB_SEARCH,
|
||||
{ path, type: GlobSearchType.DIR, include, max_depth: maxDepth, limit: SEARCH_LIMIT },
|
||||
signal
|
||||
);
|
||||
if (typeof res.error === 'string') return { base: '', entries: [], error: res.error };
|
||||
const base = typeof res.base === 'string' ? res.base : '';
|
||||
const entries = Array.isArray(res.entries) ? (res.entries as GlobEntry[]) : [];
|
||||
searchCache.set(key, { results: entries, base, at: Date.now() });
|
||||
return { base, entries };
|
||||
}
|
||||
|
||||
async function doSearch(query: string) {
|
||||
const trimmed = query.trim();
|
||||
if (!trimmed) {
|
||||
queryResults = [];
|
||||
searchError = null;
|
||||
isSearching = false;
|
||||
hoveredIndex = -1;
|
||||
searchScope = homeBase ?? HOME_TILDE;
|
||||
return;
|
||||
}
|
||||
|
||||
cancelSearch();
|
||||
const controller = new AbortController();
|
||||
searchController = controller;
|
||||
const mySeq = ++searchSeq;
|
||||
|
||||
const pathQuery = splitPathQuery(trimmed);
|
||||
|
||||
isSearching = true;
|
||||
try {
|
||||
// A generous limit is requested because ranking happens
|
||||
// client-side; only the top 20 are shown.
|
||||
const searchPath = pathQuery ? pathQuery.parent : (homeBase ?? HOME_TILDE);
|
||||
const include = pathQuery
|
||||
? pathQuery.last
|
||||
? buildCaseInsensitiveGlob(pathQuery.last)
|
||||
: GLOB_WILDCARD
|
||||
: buildCaseInsensitiveGlob(trimmed);
|
||||
const maxDepth = pathQuery ? PATH_NAV_MAX_DEPTH : SEARCH_MAX_DEPTH;
|
||||
const res = await searchDirs(searchPath, include, maxDepth, controller.signal);
|
||||
if (mySeq !== searchSeq) return;
|
||||
if (res.error) {
|
||||
queryResults = [];
|
||||
hoveredIndex = -1;
|
||||
searchError = res.error;
|
||||
return;
|
||||
}
|
||||
const { base, entries } = res;
|
||||
const ranked = rankEntries(entries, pathQuery?.last ?? trimmed);
|
||||
let results = ranked.map((e) => joinPath(base, e.path));
|
||||
searchScope = pathQuery ? pathQuery.parent : (homeBase ?? HOME_TILDE);
|
||||
|
||||
// An exactly-typed directory is "entered": list its children too,
|
||||
// so path navigation doesn't require a trailing slash.
|
||||
const last = pathQuery?.last;
|
||||
const exact = last
|
||||
? ranked.find((e) => lastPathSegment(e.path).toLowerCase() === last.toLowerCase())
|
||||
: undefined;
|
||||
if (exact) {
|
||||
const exactDir = joinPath(base, exact.path);
|
||||
const childRes = await searchDirs(
|
||||
exactDir,
|
||||
GLOB_WILDCARD,
|
||||
PATH_NAV_MAX_DEPTH,
|
||||
controller.signal
|
||||
);
|
||||
if (mySeq !== searchSeq) return;
|
||||
if (!childRes.error) {
|
||||
const children = childRes.entries
|
||||
.map((e) => joinPath(childRes.base, e.path))
|
||||
.sort((a, b) => a.localeCompare(b));
|
||||
results = [...results, ...children];
|
||||
searchScope = exactDir;
|
||||
}
|
||||
}
|
||||
|
||||
queryResults = results.slice(0, MAX_RESULTS_SHOWN);
|
||||
hoveredIndex = queryResults.length > 0 ? 0 : -1;
|
||||
// new results: scroll the list back to the top (first item is hovered)
|
||||
if (hoveredIndex === 0) scrollTrigger++;
|
||||
searchError = null;
|
||||
} catch (err) {
|
||||
if (mySeq !== searchSeq) return;
|
||||
queryResults = [];
|
||||
hoveredIndex = -1;
|
||||
if (controller.signal.aborted) return;
|
||||
searchError = err instanceof Error ? err.message : String(err);
|
||||
} finally {
|
||||
if (mySeq === searchSeq) isSearching = false;
|
||||
}
|
||||
}
|
||||
|
||||
// Single funnel for every local close so the host refocus fires
|
||||
// regardless of which commit/dismiss path ended the interaction.
|
||||
function closePicker() {
|
||||
isOpen = false;
|
||||
onClose?.();
|
||||
}
|
||||
|
||||
function commit(path: string) {
|
||||
directory = path;
|
||||
onChange?.(path);
|
||||
closePicker();
|
||||
}
|
||||
|
||||
function setDirectory(value: string) {
|
||||
const trimmed = value.trim();
|
||||
if (!trimmed) return;
|
||||
directory = trimmed;
|
||||
onChange?.(trimmed);
|
||||
}
|
||||
|
||||
// Resolve a folder name picked via the browser-native picker (which exposes
|
||||
// only the leaf name) to a server-side absolute path. Returns null when the
|
||||
// server cannot locate a matching directory, so the caller can fail visibly
|
||||
// instead of committing a bare leaf name that would resolve against the
|
||||
// server process working directory.
|
||||
async function resolveNativeName(name: string): Promise<string | null> {
|
||||
try {
|
||||
const res = await ToolsService.executeToolRaw(BuiltInTool.FILE_GLOB_SEARCH, {
|
||||
path: homeBase ?? HOME_TILDE,
|
||||
type: GlobSearchType.DIR,
|
||||
include: buildCaseInsensitiveGlob(name),
|
||||
max_depth: NATIVE_MAX_DEPTH,
|
||||
limit: NATIVE_LIMIT
|
||||
});
|
||||
const base = typeof res.base === 'string' ? res.base : '';
|
||||
const entries = Array.isArray(res.entries) ? (res.entries as GlobEntry[]) : [];
|
||||
const match = entries.find(
|
||||
(e) => lastPathSegment(e.path).toLowerCase() === name.toLowerCase()
|
||||
);
|
||||
return match ? joinPath(base, match.path) : null;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
async function browseNative() {
|
||||
if (disabled || !window.showDirectoryPicker) return;
|
||||
try {
|
||||
const handle = await window.showDirectoryPicker();
|
||||
const path = await resolveNativeName(handle.name);
|
||||
if (path) {
|
||||
setDirectory(path);
|
||||
closePicker();
|
||||
} else {
|
||||
// keep the previous cwd and fail visibly instead of committing a
|
||||
// bare leaf name that would resolve against the server cwd
|
||||
searchError = `Could not resolve "${handle.name}" to a server path`;
|
||||
}
|
||||
} catch (err) {
|
||||
// user cancelled - silently ignore; other errors are logged
|
||||
if (err instanceof DOMException && err.name === 'AbortError') return;
|
||||
console.error('[ChatFormWorkingDirectory] showDirectoryPicker failed:', err);
|
||||
}
|
||||
}
|
||||
|
||||
function handleSubmit() {
|
||||
const value = inputValue.trim();
|
||||
if (!value) {
|
||||
closePicker();
|
||||
return;
|
||||
}
|
||||
setDirectory(value);
|
||||
closePicker();
|
||||
}
|
||||
|
||||
function handleKeydown(event: KeyboardEvent) {
|
||||
if (event.key === KeyboardKey.ENTER) {
|
||||
event.preventDefault();
|
||||
// Commit the highlighted result, falling back to the raw input
|
||||
// only when the query returned no matches.
|
||||
if (hoveredIndex >= 0 && queryResults[hoveredIndex]) {
|
||||
commit(queryResults[hoveredIndex]);
|
||||
} else if (queryResults.length === 0) {
|
||||
handleSubmit();
|
||||
}
|
||||
} else if (event.key === KeyboardKey.ARROW_DOWN) {
|
||||
if (queryResults.length > 0) {
|
||||
event.preventDefault();
|
||||
hoveredIndex = (hoveredIndex + 1) % queryResults.length;
|
||||
scrollTrigger++;
|
||||
}
|
||||
} else if (event.key === KeyboardKey.ARROW_UP) {
|
||||
if (queryResults.length > 0) {
|
||||
event.preventDefault();
|
||||
hoveredIndex = hoveredIndex <= 0 ? queryResults.length - 1 : hoveredIndex - 1;
|
||||
scrollTrigger++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function handleInputInput(value: string) {
|
||||
hoveredIndex = -1;
|
||||
if (value.trim().length > 0) {
|
||||
runSearch(value);
|
||||
}
|
||||
}
|
||||
|
||||
function clearDirectory(event?: MouseEvent) {
|
||||
// Stop the click from bubbling into the popover trigger and re-opening
|
||||
// the picker on top of the now-cleared state.
|
||||
event?.stopPropagation();
|
||||
event?.preventDefault();
|
||||
directory = null;
|
||||
onChange?.(null);
|
||||
closePicker();
|
||||
}
|
||||
|
||||
// The chip is always visible; the X clears the directory (no-op when
|
||||
// already empty).
|
||||
function handleDismiss(event?: MouseEvent) {
|
||||
event?.stopPropagation();
|
||||
event?.preventDefault();
|
||||
if (directory) {
|
||||
clearDirectory(event);
|
||||
}
|
||||
}
|
||||
|
||||
function handleOpenChange(open: boolean) {
|
||||
isOpen = open;
|
||||
if (open) {
|
||||
// Seed the search field with the current path so the user can refine it
|
||||
// (or hit Enter to confirm / clear via the X icon).
|
||||
inputValue = directory ?? '';
|
||||
hoveredIndex = -1;
|
||||
queryResults = [];
|
||||
searchError = null;
|
||||
void toolsStore.resolveServerHome();
|
||||
searchScope = homeBase ?? HOME_TILDE;
|
||||
if (inputValue.trim()) runSearch(inputValue);
|
||||
} else {
|
||||
cancelSearch();
|
||||
// bits-ui-initiated close (Escape on the content, outside-click,
|
||||
// trigger toggle) - the only path that bypasses closePicker().
|
||||
onClose?.();
|
||||
}
|
||||
}
|
||||
|
||||
// Tooltips only on wider viewports - hover surfaces get in the way on
|
||||
// touch / narrow layouts. Mirrors the gate used in ActionIcon.
|
||||
let innerWidth = $state(0);
|
||||
const showTooltip = $derived(innerWidth > DEFAULT_MOBILE_BREAKPOINT);
|
||||
</script>
|
||||
|
||||
<div
|
||||
class={[
|
||||
'justify-self-start flex min-w-0 w-auto items-center gap-1 mt-1.5 py-1 px-2 backdrop-blur-2xl rounded-md',
|
||||
className,
|
||||
isOpen && 'w-full'
|
||||
]}
|
||||
>
|
||||
<Popover.Root bind:open={isOpen} onOpenChange={handleOpenChange}>
|
||||
<Popover.Trigger {disabled} class="flex justify-start">
|
||||
<ChatFormWorkingDirectoryChip
|
||||
{directory}
|
||||
{homeBase}
|
||||
{disabled}
|
||||
{showTooltip}
|
||||
onClear={handleDismiss}
|
||||
/>
|
||||
</Popover.Trigger>
|
||||
|
||||
<Popover.Content
|
||||
side="top"
|
||||
align="start"
|
||||
sideOffset={4}
|
||||
class="md:max-w-3xl w-[calc(100vw-1rem)] rounded-xl border-border/50 p-0 shadow-xl md:-translate-2!"
|
||||
onkeydown={handleKeydown}
|
||||
onOpenAutoFocus={(event) => event.preventDefault()}
|
||||
>
|
||||
<div class="p-2 min-h-28 flex flex-col justify-between">
|
||||
<SearchInput
|
||||
bind:ref={searchInputRef}
|
||||
bind:value={inputValue}
|
||||
placeholder="Choose working directory"
|
||||
onInput={handleInputInput}
|
||||
onClose={closePicker}
|
||||
class="w-full"
|
||||
/>
|
||||
|
||||
{#if inputValue.trim() && (isSearching || queryResults.length > 0 || searchError)}
|
||||
<ChatFormWorkingDirectoryResultsList
|
||||
results={queryResults}
|
||||
{hoveredIndex}
|
||||
{isSearching}
|
||||
error={searchError}
|
||||
rawQuery={inputValue}
|
||||
bind:container={listContainer}
|
||||
onCommit={commit}
|
||||
onHover={(index) => (hoveredIndex = index)}
|
||||
/>
|
||||
{/if}
|
||||
|
||||
{#if pickerSupported}
|
||||
<button
|
||||
type="button"
|
||||
class="-mt-1 flex cursor-pointer items-center gap-2 rounded-sm px-2 py-1.5 text-sm outline-hidden select-none hover:bg-accent hover:text-accent-foreground"
|
||||
onclick={browseNative}
|
||||
>
|
||||
<FolderOpen class="size-4 shrink-0 text-muted-foreground" />
|
||||
<span>Browse</span>
|
||||
</button>
|
||||
{/if}
|
||||
|
||||
{#if homeBase}
|
||||
<div class="-mx-2 my-1 h-px bg-border/20" aria-hidden="true"></div>
|
||||
|
||||
<span class="px-2 py-2 font-mono text-[10px]">
|
||||
Searching in:
|
||||
|
||||
<span class="truncate text-muted-foreground/70" title={searchScope}
|
||||
>{abbreviateHome(searchScope, homeBase)}</span
|
||||
>
|
||||
</span>
|
||||
{/if}
|
||||
</div>
|
||||
</Popover.Content>
|
||||
</Popover.Root>
|
||||
</div>
|
||||
|
||||
<svelte:window bind:innerWidth />
|
||||
@@ -0,0 +1,69 @@
|
||||
<script lang="ts">
|
||||
import { Folder, X } from '@lucide/svelte';
|
||||
import { abbreviateWorkingDir } from '$lib/utils';
|
||||
import * as Tooltip from '$lib/components/ui/tooltip';
|
||||
import { ActionIcon } from '$lib/components/app/actions';
|
||||
|
||||
interface Props {
|
||||
directory?: string | null;
|
||||
homeBase?: string | null;
|
||||
disabled?: boolean;
|
||||
showTooltip?: boolean;
|
||||
onClear?: (event?: MouseEvent) => void;
|
||||
}
|
||||
|
||||
let {
|
||||
directory = null,
|
||||
homeBase = null,
|
||||
disabled = false,
|
||||
showTooltip = false,
|
||||
onClear
|
||||
}: Props = $props();
|
||||
|
||||
const displayLabel = $derived(
|
||||
directory ? abbreviateWorkingDir(directory, homeBase) : 'Select working directory'
|
||||
);
|
||||
// Full path surface: hover the abbreviated label to recall the exact directory.
|
||||
const displayLabelTitle = $derived(directory ?? '');
|
||||
</script>
|
||||
|
||||
<span
|
||||
class="text-muted-foreground inline-flex items-center gap-1 text-xs group"
|
||||
class:text-foreground={directory}
|
||||
>
|
||||
<div class="flex min-w-0 items-center gap-1 cursor-pointer">
|
||||
<Folder class="w-3.5 h-3.5" />
|
||||
|
||||
{#if showTooltip && displayLabelTitle}
|
||||
<Tooltip.Root>
|
||||
<Tooltip.Trigger>
|
||||
{#snippet child({ props })}
|
||||
<span {...props} class="max-w-64 truncate">{displayLabel}</span>
|
||||
{/snippet}
|
||||
</Tooltip.Trigger>
|
||||
<Tooltip.Content>
|
||||
<p>{displayLabelTitle}</p>
|
||||
</Tooltip.Content>
|
||||
</Tooltip.Root>
|
||||
{:else}
|
||||
<span class="max-w-64 truncate">{displayLabel}</span>
|
||||
{/if}
|
||||
</div>
|
||||
|
||||
{#if directory}
|
||||
<div
|
||||
class="w-0 overflow-hidden opacity-0 transition-[width,opacity] duration-200 ease-out group-hover:w-auto group-hover:opacity-100"
|
||||
>
|
||||
<ActionIcon
|
||||
icon={X}
|
||||
tooltip="Reset working directory"
|
||||
ariaLabel="Reset working directory"
|
||||
{disabled}
|
||||
onclick={onClear}
|
||||
iconSize="h-3 w-3"
|
||||
stopPropagationOnClick
|
||||
class="!h-4 !w-4 shrink-0 text-muted-foreground hover:text-foreground"
|
||||
/>
|
||||
</div>
|
||||
{/if}
|
||||
</span>
|
||||
+72
@@ -0,0 +1,72 @@
|
||||
<script lang="ts">
|
||||
import { Folder } from '@lucide/svelte';
|
||||
import { fly } from 'svelte/transition';
|
||||
import { highlightMatch } from '$lib/utils';
|
||||
import { cn } from '$lib/components/ui/utils';
|
||||
|
||||
// Fly-in transition for the results list.
|
||||
const FLY_Y_PX = -4;
|
||||
const FLY_DURATION_MS = 100;
|
||||
|
||||
interface Props {
|
||||
results: string[];
|
||||
hoveredIndex: number;
|
||||
isSearching: boolean;
|
||||
error: string | null;
|
||||
rawQuery: string;
|
||||
container?: HTMLDivElement | null;
|
||||
onCommit?: (path: string) => void;
|
||||
onHover?: (index: number) => void;
|
||||
}
|
||||
|
||||
let {
|
||||
results,
|
||||
hoveredIndex,
|
||||
isSearching,
|
||||
error,
|
||||
rawQuery,
|
||||
container = $bindable(null),
|
||||
onCommit,
|
||||
onHover
|
||||
}: Props = $props();
|
||||
</script>
|
||||
|
||||
<div
|
||||
bind:this={container}
|
||||
class="max-h-48 overflow-y-auto py-2"
|
||||
transition:fly={{ y: FLY_Y_PX, duration: FLY_DURATION_MS }}
|
||||
>
|
||||
{#if isSearching && results.length === 0}
|
||||
<div class="px-2 py-1.5 text-sm text-muted-foreground">Searching...</div>
|
||||
{:else if error}
|
||||
<div class="px-2 py-1.5 text-sm text-destructive">{error}</div>
|
||||
{:else if results.length === 0}
|
||||
<div class="px-2 py-1.5 text-sm text-muted-foreground">No matching folders</div>
|
||||
{:else}
|
||||
{#each results as path, index (path)}
|
||||
<button
|
||||
type="button"
|
||||
data-result-index={index}
|
||||
data-highlighted={index === hoveredIndex ? '' : undefined}
|
||||
class={cn(
|
||||
'relative flex w-full cursor-pointer items-center gap-2 rounded-sm px-2 py-1.5 text-sm outline-hidden select-none data-highlighted:bg-accent data-highlighted:text-accent-foreground'
|
||||
)}
|
||||
onclick={() => onCommit?.(path)}
|
||||
onmouseenter={() => onHover?.(index)}
|
||||
>
|
||||
<Folder class="size-4 shrink-0 text-muted-foreground" />
|
||||
<span class="min-w-0 flex-1 truncate font-mono text-left">
|
||||
{#each highlightMatch(path, rawQuery.trim()) as seg, segIndex (segIndex)}
|
||||
{#if seg.match}
|
||||
<mark class="rounded bg-yellow-200/60 px-0.5 text-foreground dark:bg-yellow-500/30"
|
||||
>{seg.text}</mark
|
||||
>
|
||||
{:else}
|
||||
{seg.text}
|
||||
{/if}
|
||||
{/each}
|
||||
</span>
|
||||
</button>
|
||||
{/each}
|
||||
{/if}
|
||||
</div>
|
||||
@@ -12,6 +12,7 @@
|
||||
ChatMessageAssistant,
|
||||
ChatMessageUser,
|
||||
ChatMessageSystem,
|
||||
ChatMessageSynthetic,
|
||||
ChatMessageMcpPrompt
|
||||
} from '$lib/components/app/chat';
|
||||
import { parseFilesToMessageExtras } from '$lib/utils/browser-only';
|
||||
@@ -56,6 +57,10 @@
|
||||
: message.content
|
||||
);
|
||||
|
||||
// Synthetic cwd-change messages render with the folder-row UI instead
|
||||
// of a user bubble. The persisted flag is the single source of truth.
|
||||
let isSynthetic = $derived(Boolean(message.isSynthetic));
|
||||
|
||||
let rawEditContent = $derived.by(() => {
|
||||
if (message.role !== MessageRole.ASSISTANT) return undefined;
|
||||
|
||||
@@ -344,7 +349,7 @@
|
||||
}
|
||||
</script>
|
||||
|
||||
<div class="chat-message">
|
||||
<div class="chat-message" class:chat-message--synthetic={isSynthetic}>
|
||||
{#if message.role === MessageRole.SYSTEM}
|
||||
<ChatMessageSystem
|
||||
bind:textareaElement
|
||||
@@ -375,6 +380,8 @@
|
||||
{showDeleteDialog}
|
||||
{siblingInfo}
|
||||
/>
|
||||
{:else if isSynthetic}
|
||||
<ChatMessageSynthetic {message} class={className} />
|
||||
{:else if message.role === MessageRole.USER}
|
||||
<ChatMessageUser
|
||||
class={className}
|
||||
@@ -422,7 +429,17 @@
|
||||
* once known; 500px sizes messages that have never been rendered.
|
||||
*/
|
||||
.chat-message {
|
||||
--chat-message-intrinsic-size: 500px;
|
||||
content-visibility: auto;
|
||||
contain-intrinsic-size: auto 500px;
|
||||
contain-intrinsic-size: auto var(--chat-message-intrinsic-size);
|
||||
}
|
||||
|
||||
/*
|
||||
* Synthetic rows (e.g. the working-directory change) are small, so an
|
||||
* accurate placeholder keeps the injected row from inflating the
|
||||
* auto-scroll offset; the 500px default is for ordinary bubbles.
|
||||
*/
|
||||
.chat-message--synthetic {
|
||||
--chat-message-intrinsic-size: 40px;
|
||||
}
|
||||
</style>
|
||||
|
||||
+31
@@ -0,0 +1,31 @@
|
||||
<script lang="ts">
|
||||
import { Folder, FolderX } from '@lucide/svelte';
|
||||
import { parseCwdMessage } from '$lib/utils';
|
||||
import type { DatabaseMessage } from '$lib/types';
|
||||
|
||||
interface Props {
|
||||
class?: string;
|
||||
message: DatabaseMessage;
|
||||
}
|
||||
|
||||
let { class: className = '', message }: Props = $props();
|
||||
|
||||
// Parse the synthetic message content in the UI so the row reuses the
|
||||
// exact same text the model saw, including any guidance suffix.
|
||||
let info = $derived(parseCwdMessage(message.content));
|
||||
</script>
|
||||
|
||||
{#if info}
|
||||
<div class="text-muted-foreground flex items-center gap-2 py-1.5 {className}">
|
||||
{#if info.path === null}
|
||||
<FolderX class="text-muted-foreground/60 h-3.5 w-3.5 shrink-0" />
|
||||
<span class="text-foreground/80 text-sm font-medium">Working directory cleared</span>
|
||||
{:else}
|
||||
<Folder class="text-muted-foreground/60 h-3.5 w-3.5 shrink-0" />
|
||||
<span class="text-foreground/80 text-sm font-medium">Set working directory to </span>
|
||||
<span class="font-mono text-foreground/90 text-sm break-all" title={info.path}>
|
||||
{info.display}
|
||||
</span>
|
||||
{/if}
|
||||
</div>
|
||||
{/if}
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
<script lang="ts">
|
||||
import { parseCwdMessage } from '$lib/utils';
|
||||
import type { DatabaseMessage } from '$lib/types';
|
||||
import ChatMessageCwdChange from './ChatMessageCwdChange.svelte';
|
||||
|
||||
interface Props {
|
||||
class?: string;
|
||||
message: DatabaseMessage;
|
||||
}
|
||||
|
||||
let { class: className = '', message }: Props = $props();
|
||||
|
||||
// Synthetic messages render a dedicated UI, never a user bubble. The only
|
||||
// kind today is the working-directory change; parse the content so the
|
||||
// row reuses the exact synthetic text (and future kinds slot in here).
|
||||
let isCwdChange = $derived(parseCwdMessage(message.content) !== null);
|
||||
</script>
|
||||
|
||||
{#if isCwdChange}
|
||||
<ChatMessageCwdChange {message} class={className} />
|
||||
{:else}
|
||||
<span class="text-muted-foreground block text-sm {className}">{message.content}</span>
|
||||
{/if}
|
||||
+3
@@ -12,6 +12,7 @@
|
||||
import ChatMessageToolCallBlockExecShellCommand from './ChatMessageToolCallBlockExecShellCommand.svelte';
|
||||
import ChatMessageToolCallBlockFileGlobSearch from './ChatMessageToolCallBlockFileGlobSearch.svelte';
|
||||
import ChatMessageToolCallBlockGetDatetime from './ChatMessageToolCallBlockGetDatetime.svelte';
|
||||
import ChatMessageToolCallBlockGetInfo from './ChatMessageToolCallBlockGetInfo.svelte';
|
||||
import ChatMessageToolCallBlockGrepSearch from './ChatMessageToolCallBlockGrepSearch.svelte';
|
||||
import ChatMessageToolCallBlockReadFile from './ChatMessageToolCallBlockReadFile.svelte';
|
||||
import ChatMessageToolCallBlockRunJavascript from './ChatMessageToolCallBlockRunJavascript.svelte';
|
||||
@@ -40,6 +41,8 @@
|
||||
<ChatMessageToolCallBlockSearchResults {section} {open} {isStreaming} {onToggle} />
|
||||
{:else if section.toolName === BuiltInTool.GET_DATETIME}
|
||||
<ChatMessageToolCallBlockGetDatetime {section} {isStreaming} />
|
||||
{:else if section.toolName === BuiltInTool.GET_INFO}
|
||||
<ChatMessageToolCallBlockGetInfo {section} {isStreaming} />
|
||||
{:else if section.toolName === BuiltInTool.READ_FILE}
|
||||
<ChatMessageToolCallBlockReadFile {section} {open} {isStreaming} {onToggle} />
|
||||
{:else if section.toolName === BuiltInTool.EDIT_FILE}
|
||||
|
||||
+6
-2
@@ -1,7 +1,8 @@
|
||||
<script lang="ts">
|
||||
import { XCircle } from '@lucide/svelte';
|
||||
import { MAX_HEIGHT_CODE_BLOCK, RESULT_STAT_SEPARATOR } from '$lib/constants';
|
||||
import { computeLineDiff, prefixFor, type AgenticSection } from '$lib/utils';
|
||||
import { computeLineDiff, prefixFor, abbreviateHome, type AgenticSection } from '$lib/utils';
|
||||
import { toolsStore } from '$lib/stores/tools.svelte';
|
||||
import { parseEditFileMeta } from './parsers/edit-file';
|
||||
import ToolCallBlock from './ToolCallBlock.svelte';
|
||||
|
||||
@@ -15,6 +16,7 @@
|
||||
let { section, open, isStreaming, onToggle }: Props = $props();
|
||||
|
||||
const editFileMeta = $derived(parseEditFileMeta(section));
|
||||
const home = $derived(toolsStore.serverHome);
|
||||
const editDiffs = $derived(
|
||||
(editFileMeta?.edits ?? []).map((edit) => computeLineDiff(edit.oldText, edit.newText))
|
||||
);
|
||||
@@ -23,7 +25,9 @@
|
||||
<ToolCallBlock {section} {open} {isStreaming} meta={editFileMeta} {onToggle}>
|
||||
{#snippet titleSnippet()}
|
||||
<span class="text-muted-foreground">Edit file </span>
|
||||
<span class="font-mono">{editFileMeta?.filePath}</span>
|
||||
<span class="font-mono" title={editFileMeta?.filePath}
|
||||
>{abbreviateHome(editFileMeta?.filePath ?? '', home)}</span
|
||||
>
|
||||
{#if editFileMeta?.errorMessage}
|
||||
<span class="ml-1 text-xs italic text-muted-foreground/70">(failed)</span>
|
||||
{/if}
|
||||
|
||||
+32
@@ -12,6 +12,7 @@
|
||||
import { config } from '$lib/stores/settings.svelte';
|
||||
import { TOOL_RUNTIME_SCROLL_AT_BOTTOM_THRESHOLD_PX } from '$lib/constants/auto-scroll';
|
||||
import {
|
||||
abbreviateHome,
|
||||
highlightCode,
|
||||
isExitCodeSummaryLine,
|
||||
parseExecShellCommandError,
|
||||
@@ -21,6 +22,7 @@
|
||||
type ExecShellExitStatus,
|
||||
type ToolResultLine
|
||||
} from '$lib/utils';
|
||||
import { toolsStore } from '$lib/stores/tools.svelte';
|
||||
import { parseExecShellCommandMeta } from './parsers/exec-shell-command';
|
||||
import type { DatabaseMessageExtra } from '$lib/types';
|
||||
import ToolCallBlock from './ToolCallBlock.svelte';
|
||||
@@ -75,6 +77,14 @@
|
||||
execShellMeta ? highlightCode(execShellMeta.command, 'bash') : ''
|
||||
);
|
||||
|
||||
// The working directory the command ran with, persisted per call on the
|
||||
// tool result message (it travels via the x-tool-cwd header, not the tool
|
||||
// args). Reading it from the section keeps it accurate even if the
|
||||
// conversation cwd changes later.
|
||||
const cwd = $derived(section.toolCwd);
|
||||
const home = $derived(toolsStore.serverHome);
|
||||
const wdDisplay = $derived(abbreviateHome(cwd ?? '', home));
|
||||
|
||||
const exitBadgeClass = $derived(
|
||||
execShellExitStatus?.timedOut
|
||||
? 'exit-badge warning'
|
||||
@@ -159,6 +169,11 @@
|
||||
</script>
|
||||
|
||||
{#snippet execShellTitle()}
|
||||
{#if cwd}
|
||||
<span class="exec-wd" title={cwd}>{wdDisplay}</span>
|
||||
<span class="exec-prompt">$</span>
|
||||
{/if}
|
||||
|
||||
{#if highlightedCommandHtml}
|
||||
<span class="font-mono">{@html highlightedCommandHtml}</span>
|
||||
{:else}
|
||||
@@ -232,6 +247,23 @@
|
||||
</ToolCallBlock>
|
||||
|
||||
<style>
|
||||
:root {
|
||||
--exec-wd-margin: 0.4rem;
|
||||
}
|
||||
|
||||
.exec-wd {
|
||||
font-family: var(--font-mono);
|
||||
color: var(--muted-foreground);
|
||||
margin-right: var(--exec-wd-margin);
|
||||
}
|
||||
|
||||
.exec-prompt {
|
||||
font-family: var(--font-mono);
|
||||
color: var(--muted-foreground);
|
||||
opacity: 0.55;
|
||||
margin-right: var(--exec-wd-margin);
|
||||
}
|
||||
|
||||
.terminal-output {
|
||||
overscroll-behavior: contain;
|
||||
}
|
||||
|
||||
+6
-2
@@ -1,6 +1,7 @@
|
||||
<script lang="ts">
|
||||
import { XCircle } from '@lucide/svelte';
|
||||
import { type AgenticSection } from '$lib/utils';
|
||||
import { abbreviateHome, type AgenticSection } from '$lib/utils';
|
||||
import { toolsStore } from '$lib/stores/tools.svelte';
|
||||
import { parseFileGlobSearchMeta } from './parsers/file-glob-search';
|
||||
import ToolCallBlock from './ToolCallBlock.svelte';
|
||||
|
||||
@@ -14,6 +15,7 @@
|
||||
let { section, open, isStreaming, onToggle }: Props = $props();
|
||||
|
||||
const fileGlobMeta = $derived(parseFileGlobSearchMeta(section));
|
||||
const home = $derived(toolsStore.serverHome);
|
||||
</script>
|
||||
|
||||
<ToolCallBlock {section} {open} {isStreaming} meta={fileGlobMeta} {onToggle}>
|
||||
@@ -26,7 +28,9 @@
|
||||
<span class="font-mono">{fileGlobMeta.include}</span>
|
||||
{/if}
|
||||
<span class="text-muted-foreground"> in </span>
|
||||
<span class="font-mono">{fileGlobMeta.path}</span>
|
||||
<span class="font-mono" title={fileGlobMeta.path}
|
||||
>{abbreviateHome(fileGlobMeta.path, home)}</span
|
||||
>
|
||||
{/if}
|
||||
{/snippet}
|
||||
|
||||
|
||||
+69
@@ -0,0 +1,69 @@
|
||||
<script lang="ts">
|
||||
import { Info, Loader2 } from '@lucide/svelte';
|
||||
import { AgenticSectionType } from '$lib/enums';
|
||||
import { abbreviateHome, type AgenticSection } from '$lib/utils';
|
||||
import { toolsStore } from '$lib/stores/tools.svelte';
|
||||
|
||||
interface Props {
|
||||
section: AgenticSection;
|
||||
isStreaming?: boolean;
|
||||
}
|
||||
|
||||
let { section, isStreaming = false }: Props = $props();
|
||||
|
||||
const isPending = $derived(section.type === AgenticSectionType.TOOL_CALL_PENDING);
|
||||
const isStreamingCall = $derived(section.type === AgenticSectionType.TOOL_CALL_STREAMING);
|
||||
const showSpinner = $derived(isPending || (isStreamingCall && isStreaming));
|
||||
|
||||
type GetInfoMeta = {
|
||||
os?: string;
|
||||
cwd?: string;
|
||||
errorMessage?: string;
|
||||
};
|
||||
|
||||
function parseGetInfoMeta(toolResultString: string | undefined): GetInfoMeta {
|
||||
if (!toolResultString) return {};
|
||||
|
||||
try {
|
||||
const parsed: unknown = JSON.parse(toolResultString);
|
||||
if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) {
|
||||
const obj = parsed as Record<string, unknown>;
|
||||
if (typeof obj.error === 'string') return { errorMessage: obj.error };
|
||||
return {
|
||||
os: typeof obj.os === 'string' ? obj.os : undefined,
|
||||
cwd: typeof obj.cwd === 'string' ? obj.cwd : undefined
|
||||
};
|
||||
}
|
||||
} catch {
|
||||
// not JSON - nothing to show
|
||||
}
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
const infoMeta = $derived(parseGetInfoMeta(section.toolResult));
|
||||
const home = $derived(toolsStore.serverHome);
|
||||
const cwdDisplay = $derived(abbreviateHome(infoMeta.cwd ?? '', home));
|
||||
</script>
|
||||
|
||||
<div class="text-muted-foreground flex items-center gap-2 py-1.5">
|
||||
<Info class="text-muted-foreground/60 h-3.5 w-3.5 shrink-0" />
|
||||
{#if showSpinner}
|
||||
<span class="text-foreground/80 text-sm font-medium">Runtime info</span>
|
||||
<Loader2 class="text-muted-foreground/70 h-3 w-3 animate-spin" />
|
||||
{:else if infoMeta.errorMessage}
|
||||
<span class="text-foreground/80 text-sm font-medium">Runtime info </span>
|
||||
<span class="text-red-600 text-xs italic dark:text-red-400">- {infoMeta.errorMessage}</span
|
||||
>
|
||||
{:else if infoMeta.os || infoMeta.cwd}
|
||||
<span class="text-foreground/80 text-sm font-medium">Runtime info </span>
|
||||
{#if infoMeta.os}
|
||||
<span class="font-mono text-foreground/90 text-sm">{infoMeta.os}</span>
|
||||
{/if}
|
||||
{#if infoMeta.cwd}
|
||||
<span class="font-mono text-foreground/90 text-sm" title={infoMeta.cwd}>{cwdDisplay}</span>
|
||||
{/if}
|
||||
{:else}
|
||||
<span class="text-foreground/80 text-sm font-medium">Runtime info</span>
|
||||
{/if}
|
||||
</div>
|
||||
+4
-2
@@ -1,6 +1,7 @@
|
||||
<script lang="ts">
|
||||
import { XCircle } from '@lucide/svelte';
|
||||
import { type AgenticSection } from '$lib/utils';
|
||||
import { abbreviateHome, type AgenticSection } from '$lib/utils';
|
||||
import { toolsStore } from '$lib/stores/tools.svelte';
|
||||
import { parseGrepSearchMeta } from './parsers/grep-search';
|
||||
import ToolCallBlock from './ToolCallBlock.svelte';
|
||||
|
||||
@@ -14,6 +15,7 @@
|
||||
let { section, open, isStreaming, onToggle }: Props = $props();
|
||||
|
||||
const grepMeta = $derived(parseGrepSearchMeta(section));
|
||||
const home = $derived(toolsStore.serverHome);
|
||||
</script>
|
||||
|
||||
<ToolCallBlock {section} {open} {isStreaming} meta={grepMeta} {onToggle}>
|
||||
@@ -22,7 +24,7 @@
|
||||
<span class="text-muted-foreground">Search for </span>
|
||||
<span class="font-mono">{grepMeta.pattern}</span>
|
||||
<span class="text-muted-foreground"> in </span>
|
||||
<span class="font-mono">{grepMeta.path}</span>
|
||||
<span class="font-mono" title={grepMeta.path}>{abbreviateHome(grepMeta.path, home)}</span>
|
||||
{/if}
|
||||
{/snippet}
|
||||
|
||||
|
||||
+6
-2
@@ -2,7 +2,8 @@
|
||||
import { XCircle } from '@lucide/svelte';
|
||||
import { SyntaxHighlightedCode } from '$lib/components/app';
|
||||
import { MAX_HEIGHT_CODE_BLOCK, RESULT_STAT_SEPARATOR } from '$lib/constants';
|
||||
import { type AgenticSection } from '$lib/utils';
|
||||
import { abbreviateHome, type AgenticSection } from '$lib/utils';
|
||||
import { toolsStore } from '$lib/stores/tools.svelte';
|
||||
import { parseWriteFileMeta } from './parsers/write-file';
|
||||
import ToolCallBlock from './ToolCallBlock.svelte';
|
||||
|
||||
@@ -16,12 +17,15 @@
|
||||
let { section, open, isStreaming, onToggle }: Props = $props();
|
||||
|
||||
const writeFileMeta = $derived(parseWriteFileMeta(section));
|
||||
const home = $derived(toolsStore.serverHome);
|
||||
</script>
|
||||
|
||||
<ToolCallBlock {section} {open} {isStreaming} meta={writeFileMeta} {onToggle}>
|
||||
{#snippet titleSnippet()}
|
||||
<span class="text-muted-foreground">Write file </span>
|
||||
<span class="font-mono">{writeFileMeta?.filePath}</span>
|
||||
<span class="font-mono" title={writeFileMeta?.filePath}
|
||||
>{abbreviateHome(writeFileMeta?.filePath ?? '', home)}</span
|
||||
>
|
||||
{#if writeFileMeta?.errorMessage}
|
||||
<span class="ml-1 text-xs italic text-muted-foreground/70">(failed)</span>
|
||||
{/if}
|
||||
|
||||
@@ -272,6 +272,16 @@ export { default as ChatFormMcpResourcesList } from './ChatForm/ChatFormMcpResou
|
||||
*/
|
||||
export { default as ChatFormTextarea } from './ChatForm/ChatFormTextarea.svelte';
|
||||
|
||||
/**
|
||||
* Working directory selector for agent mode. Renders a chip below the chat
|
||||
* form; clicking it opens a popover with a directory picker backed by the
|
||||
* server's `file_glob_search` built-in tool (POST /tools). The picked
|
||||
* directory is exposed via `bind:directory`; changing it records a
|
||||
* synthetic "Set working directory to ..." user message into chat history
|
||||
* and is enforced on tool calls via the `x-tool-cwd` request header.
|
||||
*/
|
||||
export { default as ChatFormWorkingDirectory } from './ChatForm/ChatFormWorkingDirectory.svelte';
|
||||
|
||||
/**
|
||||
* **ChatFormPickerMcpPrompts** - MCP prompt selection interface
|
||||
*
|
||||
@@ -557,6 +567,22 @@ export { default as ChatMessageStatisticsBadge } from './ChatMessages/ChatMessag
|
||||
*/
|
||||
export { default as ChatMessageMcpPrompt } from './ChatMessages/ChatMessage/ChatMessageMcpPrompt/ChatMessageMcpPrompt.svelte';
|
||||
|
||||
/**
|
||||
* Synthetic working-directory-change message. Rendered in place of a user
|
||||
* bubble when the message content parses as a cwd message (see
|
||||
* parseCwdMessage); shows the new cwd with the same folder-row treatment
|
||||
* the tool-call UI used.
|
||||
*/
|
||||
export { default as ChatMessageCwdChange } from './ChatMessages/ChatMessage/ChatMessageCwdChange.svelte';
|
||||
|
||||
/**
|
||||
* Generic wrapper for UI-generated (synthetic) messages. Routes the
|
||||
* working-directory change to ChatMessageCwdChange and renders a muted
|
||||
* fallback for any other synthetic text, so no synthetic message ever
|
||||
* surfaces as a user bubble.
|
||||
*/
|
||||
export { default as ChatMessageSynthetic } from './ChatMessages/ChatMessage/ChatMessageSynthetic.svelte';
|
||||
|
||||
/**
|
||||
* Formatted content display for MCP prompt messages. Renders the full prompt
|
||||
* content with arguments in a readable format. Used within ChatMessageMcpPrompt
|
||||
|
||||
@@ -15,6 +15,7 @@ import {
|
||||
FilePlus,
|
||||
FileSearch,
|
||||
FileText,
|
||||
Info,
|
||||
SearchCode,
|
||||
Terminal
|
||||
} from '@lucide/svelte';
|
||||
@@ -41,6 +42,7 @@ export const BUILTIN_TOOL_UI: Readonly<Record<BuiltInTool, BuiltinToolUiEntry>>
|
||||
source: ToolSource.BUILTIN
|
||||
},
|
||||
[BuiltInTool.GET_DATETIME]: { icon: Clock, label: 'Current time', source: ToolSource.BUILTIN },
|
||||
[BuiltInTool.GET_INFO]: { icon: Info, label: 'Runtime info', source: ToolSource.BUILTIN },
|
||||
[BuiltInTool.EXEC_SHELL_COMMAND]: {
|
||||
icon: Terminal,
|
||||
label: 'Run command',
|
||||
|
||||
@@ -40,6 +40,7 @@ export * from './mcp';
|
||||
export * from './mcp-form';
|
||||
export * from './mcp-resource';
|
||||
export * from './message-export';
|
||||
export * from './path-display';
|
||||
export * from './model-id';
|
||||
export * from './model-loading';
|
||||
export * from './sse';
|
||||
@@ -60,3 +61,4 @@ export * from './ui';
|
||||
export * from './uri-template';
|
||||
export * from './url';
|
||||
export * from './viewport';
|
||||
export * from './working-directory';
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
/**
|
||||
* Constants for synthetic working-directory messages.
|
||||
*
|
||||
* The synthetic cwd-change message is text the UI renders as a folder row
|
||||
* and the model sees as a turn reminder. The prefix and cleared marker keep
|
||||
* the human-readable wording; the file-link regexes parse the
|
||||
* `[file:///abs/path](display)` payload back out on the UI side.
|
||||
*/
|
||||
|
||||
import { UrlProtocol } from '$lib/enums';
|
||||
|
||||
export const CWD_CHANGED_PREFIX = 'Set working directory to ';
|
||||
export const CWD_CLEARED_TEXT = 'Working directory cleared';
|
||||
|
||||
export const HOME_TILDE = '~';
|
||||
export const HOME_TILDE_PREFIX = '~/'; // tilde plus path separator
|
||||
|
||||
/** Scheme prefix of the file link embedded in a synthetic cwd message. */
|
||||
export const FILE_URI_PREFIX = `${UrlProtocol.FILE}//`;
|
||||
|
||||
/** Matches the leading `[file:///abs/path](display)` link; not anchored to the end so trailing guidance may follow. */
|
||||
export const CWD_LINK_REGEX = /^\[file:\/\/([\s\S]*?)\]\(([\s\S]*?)\)/;
|
||||
@@ -1,5 +1,8 @@
|
||||
import { ToolSource } from '$lib/enums/tools.enums';
|
||||
|
||||
/** HTTP header carrying the working directory a tool call runs in. The server resolves relative paths against it; the model cannot override it. */
|
||||
export const X_TOOL_CWD_HEADER = 'x-tool-cwd';
|
||||
|
||||
export const TOOL_GROUP_LABELS = {
|
||||
[ToolSource.BUILTIN]: 'Built-in',
|
||||
[ToolSource.CUSTOM]: 'JSON Schema',
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
/**
|
||||
* Constants for the working-directory picker's glob search.
|
||||
*
|
||||
* The picker glob-matches home-relative names client-side. Character classes
|
||||
* are built case-insensitively and the reserved glob metacharacters are
|
||||
* escaped (passed through literally) so a query never changes matching.
|
||||
*/
|
||||
|
||||
export const GLOB_WILDCARD = '*';
|
||||
|
||||
/** Character that starts and ends a glob character-class fragment. */
|
||||
export const GLOB_RANGE_OPEN = '[';
|
||||
export const GLOB_RANGE_CLOSE = ']';
|
||||
|
||||
/** Query characters that carry glob meaning and are passed through literally. */
|
||||
export const GLOB_SPECIAL_CHARS = '*?[]';
|
||||
|
||||
/** Separator Windows accepts alongside `/`, and a legal POSIX filename character. */
|
||||
export const WINDOWS_SEPARATOR = '\\';
|
||||
|
||||
/** `C:`, the drive part of a Windows absolute path. */
|
||||
export const DRIVE_PREFIX_REGEX = /^[A-Za-z]:/;
|
||||
|
||||
/** `C:` or `C:/`, the root of a Windows drive-absolute path. */
|
||||
export const DRIVE_ROOT_REGEX = /^[A-Za-z]:\/?/;
|
||||
|
||||
/** `//host/share` or `//host/share/`, the root of a UNC path. */
|
||||
export const UNC_ROOT_REGEX = /^\/\/[^/]+\/[^/]+\/?/;
|
||||
|
||||
// Search tuning for the picker's file_glob_search calls.
|
||||
export const SEARCH_DEBOUNCE_MS = 180;
|
||||
export const SEARCH_LIMIT = 100;
|
||||
export const MAX_RESULTS_SHOWN = 20;
|
||||
// Home-relative globs descend deeper than path navigation, which only
|
||||
// needs the direct children of the parent.
|
||||
export const SEARCH_MAX_DEPTH = 6;
|
||||
export const PATH_NAV_MAX_DEPTH = 1;
|
||||
// Native folder-picker resolution searches a shallow, bounded window.
|
||||
export const NATIVE_MAX_DEPTH = 4;
|
||||
export const NATIVE_LIMIT = 20;
|
||||
@@ -72,6 +72,12 @@ export { ColorMode, HtmlInputType, McpPromptVariant, TooltipSide, UrlProtocol }
|
||||
|
||||
export { KeyboardKey } from './keyboard.enums';
|
||||
|
||||
export { BuiltInTool, ToolSource, ToolPermissionDecision, ToolResponseField } from './tools.enums';
|
||||
export {
|
||||
BuiltInTool,
|
||||
GlobSearchType,
|
||||
ToolSource,
|
||||
ToolPermissionDecision,
|
||||
ToolResponseField
|
||||
} from './tools.enums';
|
||||
|
||||
export { SplashOrientation } from './splash.enums';
|
||||
|
||||
@@ -17,6 +17,16 @@ export enum ToolResponseField {
|
||||
ERROR = 'error'
|
||||
}
|
||||
|
||||
/**
|
||||
* Entry types accepted by the `file_glob_search` tool's `type` parameter.
|
||||
* Mirrors the server-side validation in server-tools.cpp.
|
||||
*/
|
||||
export enum GlobSearchType {
|
||||
FILE = 'file',
|
||||
DIR = 'dir',
|
||||
ALL = 'all'
|
||||
}
|
||||
|
||||
/**
|
||||
* Wire-format identifiers for built-in and frontend tools. The string
|
||||
* value matches what the model emits in tool call names, so comparing
|
||||
@@ -30,6 +40,7 @@ export enum BuiltInTool {
|
||||
EDIT_FILE = 'edit_file',
|
||||
WRITE_FILE = 'write_file',
|
||||
GET_DATETIME = 'get_datetime',
|
||||
GET_INFO = 'get_info',
|
||||
FILE_GLOB_SEARCH = 'file_glob_search',
|
||||
GREP_SEARCH = 'grep_search',
|
||||
EXEC_SHELL_COMMAND = 'exec_shell_command',
|
||||
|
||||
@@ -24,6 +24,7 @@ export enum McpPromptVariant {
|
||||
*/
|
||||
export enum UrlProtocol {
|
||||
DATA = 'data:',
|
||||
FILE = 'file:',
|
||||
HTTP = 'http:',
|
||||
HTTPS = 'https:',
|
||||
WEBSOCKET = 'ws:',
|
||||
|
||||
@@ -674,7 +674,8 @@ export class DatabaseService {
|
||||
serverId: o.serverId,
|
||||
enabled: o.enabled
|
||||
}))
|
||||
: undefined
|
||||
: undefined,
|
||||
cwd: sourceConv.cwd
|
||||
};
|
||||
|
||||
await db[IDXDB_TABLES.conversations].add(newConv);
|
||||
|
||||
@@ -2,7 +2,7 @@ import { base } from '$app/paths';
|
||||
import { getJsonHeaders } from '$lib/utils/api-headers';
|
||||
import { parseSseJsonStream, type SseJsonEvent } from '$lib/utils/sse';
|
||||
import { apiFetch } from '$lib/utils';
|
||||
import { API_TOOLS } from '$lib/constants';
|
||||
import { API_TOOLS, X_TOOL_CWD_HEADER } from '$lib/constants';
|
||||
import { ToolResponseField } from '$lib/enums';
|
||||
import type { ToolExecutionResult, ServerBuiltinToolInfo } from '$lib/types';
|
||||
|
||||
@@ -18,15 +18,21 @@ export class ToolsService {
|
||||
|
||||
/**
|
||||
* Execute a built-in tool on the server.
|
||||
*
|
||||
* @param cwd - Working directory for the tool call, sent as the
|
||||
* x-tool-cwd request header. The server resolves relative paths
|
||||
* against it; the model cannot override it.
|
||||
*/
|
||||
static async executeTool(
|
||||
toolName: string,
|
||||
params: Record<string, unknown>,
|
||||
signal?: AbortSignal
|
||||
signal?: AbortSignal,
|
||||
cwd?: string
|
||||
): Promise<ToolExecutionResult> {
|
||||
const result = await apiFetch<Record<string, unknown>>(API_TOOLS.EXECUTE, {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ tool: toolName, params }),
|
||||
headers: cwd ? { [X_TOOL_CWD_HEADER]: cwd } : undefined,
|
||||
signal
|
||||
});
|
||||
|
||||
@@ -41,6 +47,25 @@ export class ToolsService {
|
||||
return { content: JSON.stringify(result), isError: false };
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a built-in tool and return the raw JSON response. Unlike
|
||||
* executeTool, this preserves structured fields (e.g. file_glob_search's
|
||||
* `entries` and `base`) that the flattened ToolExecutionResult drops.
|
||||
*/
|
||||
static async executeToolRaw(
|
||||
toolName: string,
|
||||
params: Record<string, unknown>,
|
||||
signal?: AbortSignal,
|
||||
cwd?: string
|
||||
): Promise<Record<string, unknown>> {
|
||||
return apiFetch<Record<string, unknown>>(API_TOOLS.EXECUTE, {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ tool: toolName, params }),
|
||||
headers: cwd ? { [X_TOOL_CWD_HEADER]: cwd } : undefined,
|
||||
signal
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Stream a built-in tool's output chunks from the server. The server
|
||||
* `POST /tools` endpoint with `{stream: true}` emits `data: {"chunk": "..."}`
|
||||
@@ -59,9 +84,11 @@ export class ToolsService {
|
||||
static async *streamTool(
|
||||
toolName: string,
|
||||
params: Record<string, unknown>,
|
||||
signal?: AbortSignal
|
||||
signal?: AbortSignal,
|
||||
cwd?: string
|
||||
): AsyncGenerator<ToolStreamEvent> {
|
||||
const headers = getJsonHeaders();
|
||||
if (cwd) headers[X_TOOL_CWD_HEADER] = cwd;
|
||||
const response = await fetch(`${base}${API_TOOLS.EXECUTE}`, {
|
||||
method: 'POST',
|
||||
headers,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user