vrr/koboldcpp
2
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-08-13 18:35:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Merge branch 'upstream' into concedo_experimental

Browse source 
# Conflicts:
#	.github/workflows/release.yml
#	.github/workflows/server.yml
#	examples/model-conversion/requirements.txt
#	examples/model-conversion/scripts/causal/run-casual-gen-embeddings-org.py
#	examples/speculative-simple/README.md
#	examples/speculative-simple/speculative-simple.cpp
#	ggml/src/ggml-opencl/ggml-opencl.cpp
#	requirements/requirements-convert_hf_to_gguf.txt
#	requirements/requirements-convert_lora_to_gguf.txt
#	scripts/hip/gcn-cdna-vgpr-check.py
#	tests/test-backend-ops.cpp
#	tests/test-chat.cpp
#	tools/imatrix/imatrix.cpp
#	tools/mtmd/CMakeLists.txt
This commit is contained in:
Concedo 2026-08-12 21:34:23 +08:00
commit f556b13a4e
79 changed files with 3714 additions and 282 deletions
Download patch file Download diff file Expand all files Collapse all files

27
common/chat.cpp
View file

@ -1180,6 +1180,16 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_
data.prompt += data.generation_prompt;
}
std::vector<std::string> tool_call_starts = { "<tool_call>" };
// Match complete <function=name> opener for Qwen3-Coder models that occasionally omit the
// starting <tool_call>. The model may hallucinate a tool name, but it is preferable over
// constraining on <function which may occur in valid content generation, e.g. #include <functional>
foreach_function(inputs.tools, [&](const json & tool) {
const std::string name = tool.at("function").at("name");
tool_call_starts.push_back("<function=" + name + ">");
});
auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
auto generation_prompt = p.literal(GEN_PREFIX);
@ -1252,7 +1262,7 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_
auto tool_calls = p.trigger_rule("tool-call-root", p.repeat(calls, min_calls, 1));
return generation_prompt +
(reasoning << p.content(p.until_one_of({ "<tool_call>", "<function=" })) << tool_calls);
(reasoning << p.content(p.until_one_of(tool_call_starts)) << tool_calls);
}
// Content only parser
@ -1278,12 +1288,9 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_
});
if (data.grammar_lazy) {
data.grammar_triggers = {
{ COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "<tool_call>" },
// Trigger on "<function" and not "<function=" because the trailing "=" is part of
// the token with the function name e.g. "=read"
{ COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "<function" },
};
for (const auto & start : tool_call_starts) {
data.grammar_triggers.push_back({ COMMON_GRAMMAR_TRIGGER_TYPE_WORD, start });
}
}
}
@ -3162,7 +3169,8 @@ static common_chat_params common_chat_params_init_muse_glimmer(const common_chat
auto analysis = p.ref("analysis");
auto recipient = p.optional(p.literal(" to=user"));
auto final_msg = p.rule("final", recipient + p.literal("<|message|>") + p.content(p.until("<|eot|>")));
auto final_msg = p.rule("final", recipient + p.literal("<|message|>") +
p.content(p.until_one_of({ "<|eot|>", "<|eom|>" })));
if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
auto string_value = p.ac(
@ -3218,7 +3226,8 @@ static common_chat_params common_chat_params_init_muse_glimmer(const common_chat
if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) {
return p.zero_or_more(start + analysis) + start + tool_calls;
}
return p.zero_or_more(start + analysis) + start + (tool_calls | final_msg);
auto trailing_calls = p.optional(p.literal("<|eom|>") + start + tool_calls);
return p.zero_or_more(start + analysis) + start + (tool_calls | (final_msg + trailing_calls));
}
return p.zero_or_more(start + analysis) + start + final_msg;

80
common/speculative.cpp
View file

@ -171,12 +171,6 @@ struct common_speculative_impl {
// (optional) serialize/restore per-seq internal state (e.g. eagle3's deferred boundary).
virtual bool get_state(llama_seq_id /*seq_id*/, std::vector<uint8_t> & /*data*/) const { return false; }
virtual void set_state(llama_seq_id /*seq_id*/, const std::vector<uint8_t> & /*data*/) {}
// true if this implementation requires the target context to extract post-norm embeddings
virtual bool need_embd() const = 0;
// true if this implementation requires the target context to extract pre-norm embeddings
virtual bool need_embd_nextn() const { return false; }
};
struct common_speculative_impl_draft_simple : public common_speculative_impl {
@ -193,6 +187,10 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
auto * ctx_dft = this->params.ctx_dft;
auto * ctx_tgt = this->params.ctx_tgt;
if (!ctx_dft) {
throw std::runtime_error("draft-simple requires a draft context");
}
SPC_TRC("%s", "adding speculative implementation 'draft-simple'\n");
SPC_TRC("- n_max=%d, n_min=%d, p_min=%f\n", this->params.n_max, this->params.n_min, this->params.p_min);
SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n",
@ -385,10 +383,6 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {
// noop
}
bool need_embd() const override {
return false;
}
};
@ -907,10 +901,6 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
pending_g_last[seq_id].resize(n_embd_dec);
std::memcpy(pending_g_last[seq_id].data(), data.data() + sizeof(llama_pos), (size_t) n_embd_dec * sizeof(float));
}
bool need_embd() const override {
return false;
}
};
// DFlash: block-diffusion drafting with a draft-side KV cache injection
@ -1247,10 +1237,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {
// noop
}
bool need_embd() const override {
return false;
}
};
struct common_speculative_impl_draft_mtp : public common_speculative_impl {
@ -1689,14 +1675,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
const size_t row_bytes = (size_t) n_embd * sizeof(float);
std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes);
}
bool need_embd() const override {
return false;
}
bool need_embd_nextn() const override {
return true;
}
};
// state of self-speculation (simple implementation, not ngram-map)
@ -1743,10 +1721,6 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl {
void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {
// noop
}
bool need_embd() const override {
return false;
}
};
struct common_speculative_impl_ngram_map_k : public common_speculative_impl {
@ -1801,10 +1775,6 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl {
common_ngram_map_accept(config[seq_id], n_accepted);
}
bool need_embd() const override {
return false;
}
};
struct common_speculative_impl_ngram_mod : public common_speculative_impl {
@ -1980,10 +1950,6 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl {
}
}
}
bool need_embd() const override {
return false;
}
};
struct common_speculative_impl_ngram_cache : public common_speculative_impl {
@ -2123,10 +2089,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {
// noop
}
bool need_embd() const override {
return false;
}
};
struct common_speculative {
@ -2322,7 +2284,6 @@ common_speculative_init_result::common_speculative_init_result(
const bool spec_mtp = std::find(params.speculative.types.begin(),
params.speculative.types.end(),
COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
GGML_ASSERT(has_draft || spec_mtp);
auto mparams = common_model_params_to_llama(params);
auto cparams = common_context_params_to_llama(params);
@ -2560,34 +2521,6 @@ bool common_speculative_process(common_speculative * spec, const llama_batch & b
return result;
}
bool common_speculative_need_embd(common_speculative * spec) {
if (spec == nullptr) {
return false;
}
for (auto & impl : spec->impls) {
if (impl->need_embd()) {
return true;
}
}
return false;
}
bool common_speculative_need_embd_nextn(common_speculative * spec) {
if (spec == nullptr) {
return false;
}
for (auto & impl : spec->impls) {
if (impl->need_embd_nextn()) {
return true;
}
}
return false;
}
void common_speculative_draft(common_speculative * spec) {
if (spec == nullptr) {
return;
@ -2672,7 +2605,10 @@ void common_speculative_draft(common_speculative * spec) {
void common_speculative_accept(common_speculative * spec, llama_seq_id seq_id, uint16_t n_accepted) {
common_speculative_impl * impl = spec->impl_last[seq_id];
GGML_ASSERT(impl);
if (impl == nullptr) {
GGML_ASSERT(n_accepted == 0);
return;
}
{
common_time_meas tm(impl->t_accept_us, !impl->gen_perf);

6
common/speculative.h
View file

@ -67,12 +67,6 @@ void common_speculative_begin(common_speculative * spec, llama_seq_id seq_id, co
// process the batch and update the internal state of the speculative context
bool common_speculative_process(common_speculative * spec, const llama_batch & batch);
// true if any implementation requires target post-norm embeddings to be extracted
bool common_speculative_need_embd(common_speculative * spec);
// true if any implementation requires target nextn embeddings to be extracted
bool common_speculative_need_embd_nextn(common_speculative * spec);
// generate drafts for the sequences specified with `common_speculative_get_draft_params`
void common_speculative_draft(common_speculative * spec);

2
conversion/__init__.py
View file

@ -214,6 +214,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"Qwen3MoeForCausalLM": "qwen",
"Qwen3NextForCausalLM": "qwen",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
"PocketTTSModel": "pockettts",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
@ -310,6 +311,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
"Qwen2_5_VLForConditionalGeneration": "qwenvl",
"Qwen3ASRForConditionalGeneration": "qwen3vl",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
"PocketTTSModel": "pockettts",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",

30
conversion/base.py
View file

@ -58,6 +58,11 @@ logger = logging.getLogger("hf-to-gguf")
AnyModel = TypeVar("AnyModel", bound="type[ModelBase]")
# for checkpoints that ship no config.json, we will try to provide a synthetic one
HparamsMatcher = Callable[[Path], bool]
HparamsLoader = Callable[[Path], dict[str, Any]]
class SentencePieceTokenTypes(IntEnum):
NORMAL = 1
UNKNOWN = 2
@ -77,6 +82,7 @@ class ModelBase:
ModelType.TEXT: {},
ModelType.MMPROJ: {},
}
_hparams_loaders: list[tuple[HparamsMatcher, HparamsLoader]] = []
dir_model: Path
ftype: gguf.LlamaFileType
@ -823,7 +829,7 @@ class ModelBase:
elif any(str(v.get("quant_algo")).endswith("NVFP4") for v in quant_layers.values() if isinstance(v, dict)):
quant_algo = "NVFP4"
self._is_nvfp4 = quant_algo == "NVFP4"
self._is_nvfp4 = quant_algo in ("NVFP4", "W4A16_NVFP4")
self._is_mxfp4 = quant_method == "mxfp4"
# NVFP4 weights are repacked and written directly to gguf_writer.
@ -1040,6 +1046,24 @@ class ModelBase:
return part_names
@staticmethod
def load_hparams_guess(dir_model: Path) -> dict[str, Any] | None:
# some models ship no config.json, will try to guess them
from conversion import load_all_models
load_all_models()
for matcher, loader in ModelBase._hparams_loaders:
if matcher(dir_model):
return loader(dir_model)
return None
@classmethod
def register_hparams_loader(cls, matcher: HparamsMatcher) -> Callable[[HparamsLoader], HparamsLoader]:
def inner(loader: HparamsLoader) -> HparamsLoader:
cls._hparams_loaders.append((matcher, loader))
return loader
return inner
@staticmethod
def load_hparams(dir_model: Path, is_mistral_format: bool):
if is_mistral_format:
@ -1053,6 +1077,10 @@ class ModelBase:
config = AutoConfig.from_pretrained(dir_model, trust_remote_code=False).to_dict()
except Exception as e:
logger.warning(f"Failed to load model config from {dir_model}: {e}")
if not (dir_model / "config.json").is_file():
config = ModelBase.load_hparams_guess(dir_model)
if config is not None:
return config
logger.warning("Trying to load config.json instead")
with open(dir_model / "config.json", "r", encoding="utf-8") as f:
config = json.load(f)

37
conversion/gemma.py
View file

@ -665,7 +665,18 @@ class Gemma4Model(Gemma3Model):
swa_layers = [t == "sliding_attention" for t in self.hparams["layer_types"]]
self.gguf_writer.add_sliding_window_pattern(swa_layers)
head_dim_full = self.hparams["global_head_dim"]
per_layer_config = self.hparams.get("per_layer_config")
layer_types = self.hparams.get("layer_types", [])
if (head_dim_full := self.hparams.get("global_head_dim")) is None and per_layer_config is not None:
for layer_idx, layer_config in per_layer_config.items():
layer_idx = int(layer_idx)
if layer_idx < len(layer_types):
if layer_types[layer_idx] == "full_attention" and "head_dim" in layer_config:
head_dim_full = layer_config["head_dim"]
break
assert head_dim_full is not None
head_dim_swa = self.hparams["head_dim"]
# correct the head dim for global/swa layers
self.gguf_writer.add_key_length(head_dim_full)
@ -685,8 +696,14 @@ class Gemma4Model(Gemma3Model):
n_ff_arr = [n_ff if il < first_kv_shared_layer_idx else n_ff * 2 for il in range(self.block_count)]
self.gguf_writer.add_feed_forward_length(n_ff_arr)
# handle num_global_key_value_heads
num_key_value_heads_full = self.hparams.get("num_global_key_value_heads")
if (num_key_value_heads_full := self.hparams.get("num_global_key_value_heads")) is None and per_layer_config is not None:
for layer_idx, layer_config in per_layer_config.items():
layer_idx = int(layer_idx)
if layer_idx < len(layer_types):
if layer_types[layer_idx] == "full_attention" and "num_key_value_heads" in layer_config:
num_key_value_heads_full = layer_config["num_key_value_heads"]
break
num_key_value_heads_swa = self.hparams.get("num_key_value_heads")
if num_key_value_heads_full is not None and num_key_value_heads_swa is not None:
value_arr = [num_key_value_heads_swa if is_swa else num_key_value_heads_full for is_swa in swa_layers]
@ -708,7 +725,19 @@ class Gemma4Model(Gemma3Model):
# IMPORTANT: this ROPE_FREQS tensor is ONLY used by the full_attention layers
rope_params_full = self.hparams["rope_parameters"]["full_attention"]
assert rope_params_full["rope_type"] == "proportional"
head_dim_full = (self.hparams["global_head_dim"])
per_layer_config = self.hparams.get("per_layer_config")
if (head_dim_full := self.hparams.get("global_head_dim")) is None and per_layer_config is not None:
layer_types = self.hparams.get("layer_types", [])
for layer_idx, layer_config in per_layer_config.items():
layer_idx = int(layer_idx)
if layer_idx < len(layer_types):
if layer_types[layer_idx] == "full_attention" and "head_dim" in layer_config:
head_dim_full = layer_config["head_dim"]
break
assert head_dim_full is not None
partial_rotary_factor_full = rope_params_full["partial_rotary_factor"]
n_rot_full = int(head_dim_full * partial_rotary_factor_full / 2)
n_unrot_full = int(head_dim_full / 2) - n_rot_full

8
conversion/nemotron.py
View file

@ -275,10 +275,18 @@ class NemotronHModel(GraniteHybridModel):
return None
elif cls.mtp_only:
# --mtp: export the MTP head plus the tensors it shares with the target model
# Include lm_head scale sidecars so NVFP4 packing sees them.
keep = name in (
"backbone.embeddings.weight",
"backbone.norm_f.weight",
"lm_head.weight",
"lm_head.weight_scale",
"lm_head.weight_scale_2",
"lm_head.weight_scale_inv",
"lm_head.input_scale",
"lm_head.input_global_scale",
"lm_head.weight_global_scale",
"lm_head.weight_packed",
)
if not keep:
return None

378
conversion/pockettts.py Normal file
View file

@ -0,0 +1,378 @@
from __future__ import annotations
import re
from pathlib import Path
from typing import Any, Iterable, TYPE_CHECKING
import torch
if TYPE_CHECKING:
from torch import Tensor
from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf, logger
# Pocket TTS is a CALM: the backbone conditions a flow-matching decoder that generates one
# continuous 32-d latent per frame. There is no codebook in this model.
# The checkpoint ships no config.json, hparams come from _load_hparams() below.
#
# Tricks being used to support this model via existing llama.cpp code paths:
# - bos_before_voice and bos_emb are learned input vectors, not tokens
# they are appended to the embedding table as extra tokens, to be looked up like any other row
# - bos_emb lives in latent space, so input_linear is folded into it here
# - the backbone has no lm_head, the embedding table is reused as output for the unused logits
#
# pipeline stage mapping:
# mimi encoder + speaker_proj --> mapped to normal mtmd audio encoder
# flow_lm.transformer --> mapped to normal libllama text model (autoregressive)
# flow_lm.flow_net + out_eos --> MTMD_GEN_PROCESS_TYPE_GEN_CODE
# mimi decoder --> MTMD_GEN_PROCESS_TYPE_GEN_WAV
# indices into mimi.encoder.model / mimi.decoder.model for stage i, see SEANetEncoder/SEANetDecoder
_ENC_RES_IDX = lambda i: 1 + 3 * i # noqa: E731
_ENC_SCALE_IDX = lambda i: 3 + 3 * i # noqa: E731
_DEC_SCALE_IDX = lambda i: 2 + 3 * i # noqa: E731
_DEC_RES_IDX = lambda i: 3 + 3 * i # noqa: E731
_N_SEANET_STAGES = 3
_SAMPLE_RATE = 24000
def _tensor_shapes(dir_model: Path) -> dict[str, tuple[int, ...]]:
part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors")
if len(part_names) != 1:
return {}
with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part:
return {name: tuple(part[name].shape) for name in part.keys()}
@ModelBase.register_hparams_loader(lambda dir_model: "flow_lm.bos_emb" in _tensor_shapes(dir_model))
def _load_hparams(dir_model: Path) -> dict[str, Any]:
logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes")
shapes = _tensor_shapes(dir_model)
n_vocab, n_embd = shapes["flow_lm.conditioner.embed.weight"]
n_layer = sum(1 for name in shapes if re.fullmatch(r"flow_lm\.transformer\.layers\.\d+\.norm1\.weight", name))
n_layer_a = sum(1 for name in shapes if re.fullmatch(r"mimi\.encoder_transformer\.transformer\.layers\.\d+\.norm1\.weight", name))
n_embd_a = shapes["mimi.encoder_transformer.transformer.layers.0.norm1.weight"][0]
return {
"architectures": ["PocketTTSModel"],
"model_type": "pockettts",
"num_hidden_layers": n_layer,
"hidden_size": n_embd,
"intermediate_size": shapes["flow_lm.transformer.layers.0.linear1.weight"][0],
# the transformer is fully causal with no context limit, this only bounds the KV cache
"max_position_embeddings": 4096,
# not in the checkpoint, but every released variant uses head_dim 64
"num_attention_heads": n_embd // 64,
# extra rows for the learned input vectors, see _embd_table()
"vocab_size": n_vocab + (2 if "flow_lm.bos_before_voice" in shapes else 1),
"rope_theta": 10000.0,
"layer_norm_eps": 1e-5,
"audio_config": {
"num_hidden_layers": n_layer_a,
"hidden_size": n_embd_a,
"intermediate_size": shapes["mimi.encoder_transformer.transformer.layers.0.linear1.weight"][0],
"num_attention_heads": n_embd_a // 64,
},
}
@ModelBase.register("PocketTTSModel")
class PocketTTSModel(TextModel):
model_arch = gguf.MODEL_ARCH.POCKETTTS
_LAYER_TENSOR_MAP = {
"norm1": gguf.MODEL_TENSOR.ATTN_NORM,
"norm2": gguf.MODEL_TENSOR.FFN_NORM,
"self_attn.out_proj": gguf.MODEL_TENSOR.ATTN_OUT,
"linear1": gguf.MODEL_TENSOR.FFN_UP,
"linear2": gguf.MODEL_TENSOR.FFN_DOWN,
}
def set_vocab(self):
# this is a unigram sentencepiece model, llama.cpp's SPM tokenizer cannot do
# unigram segmentation, so use the UGM tokenizer instead
from sentencepiece import sentencepiece_model_pb2 as model
proto = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
proto.ParseFromString(open(self.dir_model / "tokenizer.model", "rb").read())
assert proto.trainer_spec.model_type == 1, "expected a unigram tokenizer"
tokens, scores, toktypes = self._create_vocab_sentencepiece()
# the last rows of the embedding table are not sentencepiece pieces
extra = self._extra_tokens()
for i, name in enumerate(extra):
tokens[len(tokens) - len(extra) + i] = name.encode("utf-8")
toktypes[len(tokens) - len(extra) + i] = SentencePieceTokenTypes.CONTROL
scores[len(tokens) - len(extra) + i] = -1000.0
self.gguf_writer.add_tokenizer_model("t5")
self.gguf_writer.add_tokenizer_pre("default")
self.gguf_writer.add_token_list(tokens)
self.gguf_writer.add_token_scores(scores)
self.gguf_writer.add_token_types(toktypes)
self.gguf_writer.add_add_space_prefix(proto.normalizer_spec.add_dummy_prefix)
self.gguf_writer.add_remove_extra_whitespaces(proto.normalizer_spec.remove_extra_whitespaces)
if proto.normalizer_spec.precompiled_charsmap:
self.gguf_writer.add_precompiled_charsmap(proto.normalizer_spec.precompiled_charsmap)
self.gguf_writer.add_add_bos_token(False)
self.gguf_writer.add_add_eos_token(False)
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if not name.startswith("flow_lm."):
return # mimi and the flow net go to the mmproj
if name == "flow_lm.conditioner.embed.weight":
yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), self._embd_table(data_torch))
return
if name.startswith("flow_lm.out_norm."):
suffix = "." + name.rsplit(".", 1)[1]
yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT_NORM, suffix=suffix), data_torch)
return
if name.startswith("flow_lm.transformer.layers."):
assert bid is not None
key_with_suffix = name.split(f"layers.{bid}.", 1)[1]
key, suffix = key_with_suffix.rsplit(".", 1)
if key == "self_attn.in_proj":
q, k, v = data_torch.chunk(3, dim=0)
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), q)
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), k)
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), v)
return
tensor = self._LAYER_TENSOR_MAP.get(key)
if tensor is not None:
yield (self.format_tensor_name(tensor, bid, suffix="." + suffix), data_torch)
return
return
def _extra_tokens(self) -> list[str]:
# the conditioner's padding row, then the learned vectors appended by _embd_table().
# bos_before_voice only exists when the pack sets insert_bos_before_voice
names = ["<|pad|>"]
if "flow_lm.bos_before_voice" in self.model_tensors:
names.append("<|bos_before_voice|>")
names.append("<|audio_bos|>")
return names
def _embd_table(self, embed: Tensor) -> Tensor:
rows = [embed]
if "flow_lm.bos_before_voice" in self.model_tensors:
rows.append(self.model_tensors["flow_lm.bos_before_voice"]().reshape(1, -1).to(embed.dtype))
# bos_emb is a latent, it only enters the backbone through input_linear
bos_emb = self.model_tensors["flow_lm.bos_emb"]()
input_linear = self.model_tensors["flow_lm.input_linear.weight"]()
audio_bos = torch.nn.functional.linear(bos_emb.float(), input_linear.float()).reshape(1, -1)
rows.append(audio_bos.to(embed.dtype))
return torch.cat(rows, dim=0)
@ModelBase.register("PocketTTSModel")
class PocketTTSMmprojModel(MmprojModel):
has_audio_encoder = True
has_vision_encoder = False
_MIMI_TFM_MAP = {
"norm1": (gguf.MODEL_TENSOR.A_ENC_INPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM),
"norm2": (gguf.MODEL_TENSOR.A_ENC_OUTPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM),
"self_attn.out_proj": (gguf.MODEL_TENSOR.A_ENC_OUTPUT, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT),
"linear1": (gguf.MODEL_TENSOR.A_ENC_FFN_UP, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP),
"linear2": (gguf.MODEL_TENSOR.A_ENC_FFN_DOWN, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN),
"layer_scale_1.scale": (gguf.MODEL_TENSOR.A_ENC_ATTN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE),
"layer_scale_2.scale": (gguf.MODEL_TENSOR.A_ENC_FFN_SCALE_LS, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE),
}
_MIMI_TFM_QKV = (
(gguf.MODEL_TENSOR.A_ENC_ATTN_Q, gguf.MODEL_TENSOR.A_ENC_ATTN_K, gguf.MODEL_TENSOR.A_ENC_ATTN_V),
(gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V),
)
def set_gguf_parameters(self):
self.gguf_writer.add_file_type(self.ftype)
assert self.hparams_audio is not None
# voice-prompt encoder: mimi encoder + speaker_proj
self.gguf_writer.add_clip_has_audio_encoder(True)
# note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models
self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC)
self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"])
self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"])
self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"])
self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"])
self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
# mimi convolves the waveform directly, it is passed around as a 1-row "mel"
self.gguf_writer.add_audio_num_mel_bins(1)
# generation: flow-matching decoder + mimi decoder
# the SEANet and flow net hparams are constant across the family, clip.cpp holds them
self.gguf_writer.add_clip_has_gen_audio_encoder(True)
self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_GEN)
self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)
self.gguf_writer.add_gen_audio_embedding_length(self.hparams_audio["hidden_size"])
self.gguf_writer.add_gen_audio_feed_forward_length(self.hparams_audio["intermediate_size"])
self.gguf_writer.add_gen_audio_block_count(self.hparams_audio["num_hidden_layers"])
self.gguf_writer.add_gen_audio_head_count(self.hparams_audio["num_attention_heads"])
self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5)
self.gguf_writer.add_gen_audio_model_variant(self.dir_model.name)
def tensor_force_quant(self, name, new_name, bid, n_dims):
del name, bid, n_dims
# conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path
if ".seanet." in new_name or new_name in ("a.downsample.conv.weight", "a.gen.wav.upsample.weight"):
return gguf.GGMLQuantizationType.F16
return False
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
del bid # the block index of the mimi transformers is parsed here, not by the base class
T = gguf.MODEL_TENSOR
if name in ("flow_lm.bos_emb", "flow_lm.bos_before_voice", "flow_lm.conditioner.embed.weight"):
return # folded into the backbone embedding table
if name.startswith("flow_lm.transformer.") or name.startswith("flow_lm.out_norm."):
return # backbone
if name == "flow_lm.speaker_proj_weight":
yield (self.format_tensor_name(T.A_ENC_SPEAKER_PROJ), data_torch)
return
if name == "flow_lm.input_linear.weight":
yield (self.format_tensor_name(T.A_GEN_INPUT_LINEAR), data_torch)
return
if name == "flow_lm.emb_mean":
yield (self.format_tensor_name(T.A_GEN_EMB_MEAN, suffix=""), data_torch)
return
if name == "flow_lm.emb_std":
yield (self.format_tensor_name(T.A_GEN_EMB_STD, suffix=""), data_torch)
return
if name.startswith("flow_lm.out_eos."):
suffix = "." + name.rsplit(".", 1)[1]
yield (self.format_tensor_name(T.A_GEN_OUT_EOS, suffix=suffix), data_torch)
return
if name.startswith("flow_lm.flow_net."):
yield from self._flow_net_tensor(name, data_torch)
return
if name == "mimi.downsample.conv.conv.weight":
yield (self.format_tensor_name(T.A_ENC_DOWNSAMPLE_CONV), data_torch)
return
if name == "mimi.upsample.convtr.convtr.weight":
yield (self.format_tensor_name(T.A_GEN_WAV_UPSAMPLE), data_torch)
return
if name == "mimi.quantizer.output_proj.weight":
yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_OUT), data_torch.squeeze(-1))
return
if "_transformer.transformer.layers." in name:
yield from self._mimi_tfm_tensor(name, data_torch)
return
if name.startswith("mimi.encoder.model.") or name.startswith("mimi.decoder.model."):
yield from self._seanet_tensor(name, data_torch)
return
return
def _flow_net_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]:
T = gguf.MODEL_TENSOR
key = name.split("flow_lm.flow_net.", 1)[1]
suffix = "." + key.rsplit(".", 1)[1]
simple = {
"input_proj": T.A_GEN_FLOW_INPUT_PROJ,
"cond_embed": T.A_GEN_FLOW_COND_EMBD,
"final_layer.linear": T.A_GEN_FLOW_FINAL_PROJ,
"final_layer.adaLN_modulation.1": T.A_GEN_FLOW_FINAL_ADA,
}
tensor = simple.get(key.rsplit(".", 1)[0])
if tensor is not None:
yield (self.format_tensor_name(tensor, suffix=suffix), data_torch)
return
if key.startswith("time_embed."):
bid = int(key.split(".")[1])
rest = key.split(f"time_embed.{bid}.", 1)[1]
time_map = {
"freqs": (T.A_GEN_FLOW_TIME_FREQS, ""),
"mlp.0": (T.A_GEN_FLOW_TIME_UP, suffix),
"mlp.2": (T.A_GEN_FLOW_TIME_DOWN, suffix),
"mlp.3.alpha": (T.A_GEN_FLOW_TIME_NORM, ""),
}
entry = time_map.get(rest) or time_map.get(rest.rsplit(".", 1)[0])
if entry is not None:
yield (self.format_tensor_name(entry[0], bid, suffix=entry[1]), data_torch)
return
if key.startswith("res_blocks."):
bid = int(key.split(".")[1])
rest = key.split(f"res_blocks.{bid}.", 1)[1].rsplit(".", 1)[0]
blk_map = {
"in_ln": T.A_GEN_FLOW_BLK_NORM,
"mlp.0": T.A_GEN_FLOW_BLK_UP,
"mlp.2": T.A_GEN_FLOW_BLK_DOWN,
"adaLN_modulation.1": T.A_GEN_FLOW_BLK_ADA,
}
tensor = blk_map.get(rest)
if tensor is not None:
yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch)
return
def _mimi_tfm_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]:
is_decoder = name.startswith("mimi.decoder_transformer.")
bid = int(name.split("_transformer.transformer.layers.", 1)[1].split(".")[0])
key_with_suffix = name.split(f".layers.{bid}.", 1)[1]
if key_with_suffix == "self_attn.in_proj.weight":
q, k, v = data_torch.chunk(3, dim=0)
names = self._MIMI_TFM_QKV[1 if is_decoder else 0]
for tensor, part in zip(names, (q, k, v)):
yield (self.format_tensor_name(tensor, bid), part)
return
key, suffix = key_with_suffix.rsplit(".", 1)
entry = self._MIMI_TFM_MAP.get(key) or self._MIMI_TFM_MAP.get(key_with_suffix)
if entry is None:
return
tensor = entry[1 if is_decoder else 0]
suffix = ".weight" if key_with_suffix.endswith(".scale") else "." + suffix
yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch)
def _seanet_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]:
T = gguf.MODEL_TENSOR
is_decoder = name.startswith("mimi.decoder.")
idx = int(name.split(".model.", 1)[1].split(".")[0])
suffix = "." + name.rsplit(".", 1)[1]
conv_in, conv_out, res1, res2, scale = (
(T.A_GEN_WAV_SEANET_CONV_IN, T.A_GEN_WAV_SEANET_CONV_OUT, T.A_GEN_WAV_SEANET_RES_CONV1,
T.A_GEN_WAV_SEANET_RES_CONV2, T.A_GEN_WAV_SEANET_SCALE_CONV)
if is_decoder else
(T.A_ENC_SEANET_CONV_IN, T.A_ENC_SEANET_CONV_OUT, T.A_ENC_SEANET_RES_CONV1,
T.A_ENC_SEANET_RES_CONV2, T.A_ENC_SEANET_SCALE_CONV)
)
if idx == 0:
yield (self.format_tensor_name(conv_in, suffix=suffix), data_torch)
return
if idx == 3 * _N_SEANET_STAGES + 2:
yield (self.format_tensor_name(conv_out, suffix=suffix), data_torch)
return
for stage in range(_N_SEANET_STAGES):
res_idx = _DEC_RES_IDX(stage) if is_decoder else _ENC_RES_IDX(stage)
scale_idx = _DEC_SCALE_IDX(stage) if is_decoder else _ENC_SCALE_IDX(stage)
if idx == scale_idx:
yield (self.format_tensor_name(scale, stage, suffix=suffix), data_torch)
return
if idx == res_idx:
# block.1 is the dilated conv, block.3 the pointwise one (0 and 2 are ELU)
inner = int(name.split(".block.", 1)[1].split(".")[0])
tensor = res1 if inner == 1 else res2
yield (self.format_tensor_name(tensor, stage, suffix=suffix), data_torch)
return

11
conversion/qwen.py
View file

@ -647,10 +647,13 @@ class DFlashModel(Qwen3Model):
# own tokenizer logic, not the Qwen default).
from . import get_model_class
with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:
target_arch = json.load(f)["architectures"][0]
target_hparams = json.load(f)
target_arch = target_hparams["architectures"][0]
target_cls = get_model_class(target_arch)
if target_cls is not type(self):
if target_arch == "NemotronHForCausalLM":
setattr(self, "is_moe", "num_experts_per_tok" in target_hparams)
target_cls.set_vocab(self) # ty: ignore[unresolved-attribute]
else:
super().set_vocab()
@ -688,6 +691,12 @@ class DFlashModel(Qwen3Model):
name = "model." + name
return super().filter_tensors((name, gen))
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True):
return
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("Qwen3DSparkModel")
class DSparkModel(DFlashModel):

39
ggml/src/ggml-cuda/ggml-cuda.cu
View file

@ -1865,6 +1865,37 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst);
}
// returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization
// [TAG_MUL_MAT_ID_CUDA_GRAPHS]
static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int cc) {
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
return true;
}
if (dst->ne[2] <= MMVQ_MAX_BATCH_SIZE) {
if (ggml_is_quantized(src0->type)) {
if (dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)) {
return false;
}
} else if (GGML_CUDA_CC_IS_AMD(cc)) {
return false;
}
}
if (ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[2], /*n_experts=*/src0->ne[2])) {
return false;
}
if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
return false;
}
return true;
}
static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
@ -1907,7 +1938,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
}
// note: this path should not be reached when recording CUDA graphs, because it requires stream synchronization
// TODO: add asserts to verify this. should work with CUDA, HIP, etc.
GGML_ASSERT(ggml_cuda_mul_mat_id_needs_sync(dst, cc));
cudaStream_t stream = ctx.stream();
GGML_ASSERT(nb12 % nb11 == 0);
@ -2526,10 +2557,8 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) {
// [TAG_MUL_MAT_ID_CUDA_GRAPHS]
if (node->op == GGML_OP_MUL_MAT_ID) {
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
const int mmvq_mmid_max = get_mmvq_mmid_max_batch(node->src[0]->type, cc);
if (!ggml_is_quantized(node->src[0]->type) || node->ne[2] > mmvq_mmid_max) {
// under these conditions, the mul_mat_id operation will need to synchronize the stream, so we cannot use CUDA graphs
// TODO: figure out a way to enable for larger batch sizes, without hurting performance
if (ggml_cuda_mul_mat_id_needs_sync(node, cc)) {
// the mul_mat_id fallback path synchronizes the stream, so we cannot use CUDA graphs
// ref: https://github.com/ggml-org/llama.cpp/pull/18958
use_cuda_graph = false;
#ifndef NDEBUG

56
ggml/src/ggml-cuda/wkv.cu
View file

@ -141,6 +141,57 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons
}
}
template <int rows_per_block>
static __global__ void __launch_bounds__(WARP_SIZE * rows_per_block, 2)
rwkv_wkv7_f32_t1_warp_row(const int T, const int C, const int H, const float * r, const float * w, const float * k, const float * v, const float * a, const float * b, const float * s, float * dst) {
constexpr int head_size = CUDA_WKV_BLOCK_SIZE;
constexpr int half_head = head_size / 2;
const int lane = threadIdx.x;
const int row = blockIdx.y * rows_per_block + threadIdx.y;
const int bid = blockIdx.x;
const int batch_i = bid / H;
const int head_i = bid % H;
const int state_size = C * head_size;
const int head_off = head_i * head_size;
const int t = batch_i * C + head_off + row;
__shared__ float _r[head_size], _w[head_size], _k[head_size], _a[head_size], _b[head_size];
if (threadIdx.y == 0) {
_r[lane] = r[batch_i * C + head_off + lane];
_w[lane] = w[batch_i * C + head_off + lane];
_k[lane] = k[batch_i * C + head_off + lane];
_a[lane] = a[batch_i * C + head_off + lane];
_b[lane] = b[batch_i * C + head_off + lane];
_r[lane + half_head] = r[batch_i * C + head_off + lane + half_head];
_w[lane + half_head] = w[batch_i * C + head_off + lane + half_head];
_k[lane + half_head] = k[batch_i * C + head_off + lane + half_head];
_a[lane + half_head] = a[batch_i * C + head_off + lane + half_head];
_b[lane + half_head] = b[batch_i * C + head_off + lane + half_head];
}
__syncthreads();
const int64_t state_base = batch_i * state_size + head_i * head_size * head_size + row * head_size;
const float s0 = s[state_base + lane];
const float s1 = s[state_base + lane + half_head];
const float sa = warp_reduce_sum(_a[lane] * s0 + _a[lane + half_head] * s1);
const float vt = v[t];
const float st0 = s0 * _w[lane] + _k[lane] * vt + sa * _b[lane];
const float st1 = s1 * _w[lane + half_head] + _k[lane + half_head] * vt + sa * _b[lane + half_head];
const float y = warp_reduce_sum(st0 * _r[lane] + st1 * _r[lane + half_head]);
dst[T * C + state_base + lane] = st0;
dst[T * C + state_base + lane + half_head] = st1;
if (lane == 0) {
dst[t] = y;
}
}
void ggml_cuda_op_rwkv_wkv6(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const float * k_d = (const float *)dst->src[0]->data;
const float * v_d = (const float *)dst->src[1]->data;
@ -191,7 +242,10 @@ void ggml_cuda_op_rwkv_wkv7(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
GGML_ASSERT(C % H == 0);
GGML_ASSERT(C / H == CUDA_WKV_BLOCK_SIZE || C / H == CUDA_WKV_BLOCK_SIZE * 2);
if (C / H == CUDA_WKV_BLOCK_SIZE) {
if (T / B == 1 && C / H == CUDA_WKV_BLOCK_SIZE) {
constexpr int rows_per_block = 4;
rwkv_wkv7_f32_t1_warp_row<rows_per_block><<<dim3(B * H, CUDA_WKV_BLOCK_SIZE / rows_per_block), dim3(WARP_SIZE, rows_per_block), 0, stream>>>(T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d);
} else if (C / H == CUDA_WKV_BLOCK_SIZE) {
rwkv_wkv7_f32<CUDA_WKV_BLOCK_SIZE><<<B * H, C / H, 0, stream>>>(B, T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d);
} else {
rwkv_wkv7_f32<CUDA_WKV_BLOCK_SIZE * 2><<<B * H, C / H, 0, stream>>>(B, T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d);

23
ggml/src/ggml-vulkan/ggml-vulkan.cpp
View file

@ -4633,6 +4633,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1], matmul_q5_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0], matmul_q8_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K], matmul_q2_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_TQ2_0], matmul_tq2_0_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K], matmul_q3_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K], matmul_q4_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3)
CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K], matmul_q5_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3)
@ -4673,6 +4674,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
@ -4745,6 +4747,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0], matmul_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
@ -4789,6 +4792,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
@ -4879,6 +4883,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1], matmul_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0], matmul_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
@ -4927,6 +4932,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
@ -4974,6 +4980,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_q5_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_q8_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_q2_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_tq2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_q3_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_q4_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_q5_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
@ -5053,6 +5060,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f32acc, matmul_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f32acc, matmul_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0].f32acc, matmul_tq2_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f32acc, matmul_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f32acc, matmul_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f32acc, matmul_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
@ -5100,6 +5108,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
CREATE_MM(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f32acc, matmul_id_subgroup_q5_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_subgroup_q8_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_subgroup_q2_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0].f32acc, matmul_id_subgroup_tq2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_subgroup_q3_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_subgroup_q4_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_subgroup_q5_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
@ -5129,6 +5138,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
CREATE_MM(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f32acc, matmul_id_q5_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_q8_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_q2_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0].f32acc, matmul_id_tq2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_q3_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_q4_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_q5_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
@ -5232,6 +5242,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f32_f32", arr_dmmv_q5_1_f32_f32_len[reduc], arr_dmmv_q5_1_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f32_f32", arr_dmmv_q8_0_f32_f32_len[reduc], arr_dmmv_q8_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f32_f32", arr_dmmv_q2_k_f32_f32_len[reduc16], arr_dmmv_q2_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_TQ2_0][i], "mul_mat_vec_tq2_0_f32_f32", arr_dmmv_tq2_0_f32_f32_len[reduc16], arr_dmmv_tq2_0_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f32_f32", arr_dmmv_q3_k_f32_f32_len[reduc16], arr_dmmv_q3_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f32_f32", arr_dmmv_q4_k_f32_f32_len[reduc16], arr_dmmv_q4_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f32_f32", arr_dmmv_q5_k_f32_f32_len[reduc16], arr_dmmv_q5_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
@ -5259,6 +5270,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f16_f32", arr_dmmv_q5_1_f16_f32_len[reduc], arr_dmmv_q5_1_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f16_f32", arr_dmmv_q8_0_f16_f32_len[reduc], arr_dmmv_q8_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f16_f32", arr_dmmv_q2_k_f16_f32_len[reduc16], arr_dmmv_q2_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_TQ2_0][i], "mul_mat_vec_tq2_0_f16_f32", arr_dmmv_tq2_0_f16_f32_len[reduc16], arr_dmmv_tq2_0_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f16_f32", arr_dmmv_q3_k_f16_f32_len[reduc16], arr_dmmv_q3_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f16_f32", arr_dmmv_q4_k_f16_f32_len[reduc16], arr_dmmv_q4_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f16_f32", arr_dmmv_q5_k_f16_f32_len[reduc16], arr_dmmv_q5_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16);
@ -5313,6 +5325,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q5_1], "mul_mat_vec_id_q5_1_f32", arr_dmmv_id_q5_1_f32_f32_len[reduc], arr_dmmv_id_q5_1_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_f32", arr_dmmv_id_q8_0_f32_f32_len[reduc], arr_dmmv_id_q8_0_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_f32", arr_dmmv_id_q2_k_f32_f32_len[reduc16], arr_dmmv_id_q2_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_TQ2_0], "mul_mat_vec_id_tq2_0_f32", arr_dmmv_id_tq2_0_f32_f32_len[reduc16], arr_dmmv_id_tq2_0_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_f32", arr_dmmv_id_q3_k_f32_f32_len[reduc16], arr_dmmv_id_q3_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q4_K], "mul_mat_vec_id_q4_k_f32", arr_dmmv_id_q4_k_f32_f32_len[reduc16], arr_dmmv_id_q4_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16);
ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q5_K], "mul_mat_vec_id_q5_k_f32", arr_dmmv_id_q5_k_f32_f32_len[reduc16], arr_dmmv_id_q5_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16);
@ -5374,6 +5387,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ2_0], "dequant_tq2_0", dequant_tq2_0_len, dequant_tq2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_K], "dequant_q4_k", dequant_q4_k_len, dequant_q4_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_K], "dequant_q5_k", dequant_q5_k_len, dequant_q5_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1);
@ -5402,6 +5416,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_1], "get_rows_q5_1", get_rows_q5_1_len, get_rows_q5_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q8_0], "get_rows_q8_0", get_rows_q8_0_len, get_rows_q8_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q2_K], "get_rows_q2_k", get_rows_q2_k_len, get_rows_q2_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_TQ2_0], "get_rows_tq2_0", get_rows_tq2_0_len, get_rows_tq2_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q3_K], "get_rows_q3_k", get_rows_q3_k_len, get_rows_q3_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_K], "get_rows_q4_k", get_rows_q4_k_len, get_rows_q4_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_K], "get_rows_q5_k", get_rows_q5_k_len, get_rows_q5_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
@ -5430,6 +5445,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_1], "get_rows_q5_1_f32", get_rows_q5_1_f32_len, get_rows_q5_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q8_0], "get_rows_q8_0_f32", get_rows_q8_0_f32_len, get_rows_q8_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q2_K], "get_rows_q2_k_f32", get_rows_q2_k_f32_len, get_rows_q2_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_TQ2_0], "get_rows_tq2_0_f32", get_rows_tq2_0_f32_len, get_rows_tq2_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q3_K], "get_rows_q3_k_f32", get_rows_q3_k_f32_len, get_rows_q3_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_K], "get_rows_q4_k_f32", get_rows_q4_k_f32_len, get_rows_q4_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_K], "get_rows_q5_k_f32", get_rows_q5_k_f32_len, get_rows_q5_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
@ -7663,6 +7679,7 @@ static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_MXFP4:
case GGML_TYPE_NVFP4:
case GGML_TYPE_TQ2_0:
break;
default:
return nullptr;
@ -7737,6 +7754,7 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_conte
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_MXFP4:
case GGML_TYPE_NVFP4:
case GGML_TYPE_TQ2_0:
break;
default:
return nullptr;
@ -7806,6 +7824,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context *
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_MXFP4:
case GGML_TYPE_NVFP4:
case GGML_TYPE_TQ2_0:
break;
default:
return nullptr;
@ -7899,6 +7918,7 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_co
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_MXFP4:
case GGML_TYPE_NVFP4:
case GGML_TYPE_TQ2_0:
break;
default:
return nullptr;
@ -7971,6 +7991,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_MXFP4:
case GGML_TYPE_NVFP4:
case GGML_TYPE_TQ2_0:
break;
default:
return nullptr;
@ -18047,6 +18068,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_MXFP4:
case GGML_TYPE_NVFP4:
case GGML_TYPE_TQ2_0:
break;
default:
return false;
@ -18152,6 +18174,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_MXFP4:
case GGML_TYPE_NVFP4:
case GGML_TYPE_TQ2_0:
case GGML_TYPE_I32:
return true;
default:

14
ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl
View file

@ -608,6 +608,20 @@ vec2 get_dm(uint ib, uint a_offset) {
}
#endif
#if defined(DATA_A_TQ2_0)
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
// elem e -> byte qs[(e/128)*32 + e%32], bits 2*((e%128)/32); w = q - 1 (d applied via get_dm)
const uint qsi = (iqs / 128) * 32 + (iqs % 32); // iqs even -> qsi, qsi+1 in same group/level
const uint shift = 2 * ((iqs % 128) / 32);
const uvec2 qs = uvec2(data_a[a_offset + ib].qs[qsi], data_a[a_offset + ib].qs[qsi + 1]);
return vec2((qs >> shift) & 3) - 1.0;
}
vec2 get_dm(uint ib, uint a_offset) {
return vec2(float(data_a[a_offset + ib].d), 0);
}
#endif
#if defined(DATA_A_Q3_K)
vec2 dequantize(uint ib, uint iqs, uint a_offset) {
iqs /= 2;

41
ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl
View file

@ -247,6 +247,44 @@ f16vec4 dequantFuncQ8_0_v(const in decodeBufQ8_0 bl, const in uint blockCoords[2
return f16vec4(vec4(qi) * vec4(float(d)));
}
layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufTQ2_0 {
block_tq2_0 block;
};
layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufTQ2_0_packed16 {
block_tq2_0_packed16 block;
};
float16_t dequantFuncTQ2_0(const in decodeBufTQ2_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
{
decodeBufTQ2_0_packed16 bl16 = decodeBufTQ2_0_packed16(bl);
const uint idx = coordInBlock[1];
const uint qsshift = (idx & 0x60) >> 4; // 0,2,4,6
uint qs = uint32_t(bl16.block.qs[((idx & 0x80) >> 3) + ((idx & 0x1E) >> 1)]);
qs = (qs >> qsshift) & 0x0303;
qs = unpack8(qs)[idx & 1];
return bl.block.d * (float16_t(int(qs)) - float16_t(1.0));
}
f16vec4 dequantFuncTQ2_0_v(const in decodeBufTQ2_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
{
const uint idx = coordInBlock[1];
const uint qsshift = (idx & 0x60) >> 4; // 0,2,4,6
const uint qsi = ((idx & 0x80) >> 2) + (idx & 0x1C); // byte index of 4-aligned group
const uint qsw = (uint(bl.block.qs[qsi]))
| (uint(bl.block.qs[qsi + 1]) << 8)
| (uint(bl.block.qs[qsi + 2]) << 16)
| (uint(bl.block.qs[qsi + 3]) << 24);
const u8vec4 q = unpack8((qsw >> qsshift) & 0x03030303);
return bl.block.d * (f16vec4(q) - f16vec4(1.0));
}
layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufQ2_K {
block_q2_K block;
};
@ -1368,6 +1406,9 @@ f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords
#elif defined(DATA_A_Q8_0)
#define dequantFuncA dequantFuncQ8_0
#define dequantFuncA_v dequantFuncQ8_0_v
#elif defined(DATA_A_TQ2_0)
#define dequantFuncA dequantFuncTQ2_0
#define dequantFuncA_v dequantFuncTQ2_0_v
#elif defined(DATA_A_Q2_K)
#define dequantFuncA dequantFuncQ2_K
#define dequantFuncA_v dequantFuncQ2_K_v

31
ggml/src/ggml-vulkan/vulkan-shaders/dequant_tq2_0.comp Normal file
View file

@ -0,0 +1,31 @@
#version 450
#include "dequant_head.glsl"
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
void main() {
[[unroll]] for (uint wgy = 0; wgy < 256; wgy++) {
const uint i = gl_WorkGroupID.x * 256 + wgy;
if (i >= p.nel / QUANT_K) {
return;
}
const uint tid = gl_LocalInvocationID.x;
const uint ip = tid / 32; // group 0,1 (128 elems each)
const uint il = tid - 32 * ip; // byte in group 0..31
const uint y_idx = i * QUANT_K + 128 * ip + il;
const uint8_t qs = data_a[i].qs[32 * ip + il];
const FLOAT_TYPE d = FLOAT_TYPE(data_a[i].d);
data_b[y_idx + 0] = D_TYPE(d * FLOAT_TYPE(int((qs >> 0) & 3) - 1));
data_b[y_idx + 32] = D_TYPE(d * FLOAT_TYPE(int((qs >> 2) & 3) - 1));
data_b[y_idx + 64] = D_TYPE(d * FLOAT_TYPE(int((qs >> 4) & 3) - 1));
data_b[y_idx + 96] = D_TYPE(d * FLOAT_TYPE(int((qs >> 6) & 3) - 1));
}
}

102
ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_tq2_0.comp Normal file
View file

@ -0,0 +1,102 @@
#version 450
#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
#include "mul_mat_vec_base.glsl"
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
FLOAT_TYPE temp[NUM_COLS][NUM_ROWS];
// ternary TQ2_0: w = (q - 1) * d. Same qs group/level layout as q2_K, but a
// single f16 scale per 256-block and no mins:
// sum_e b_e * (q_e - 1) * d = d * (sum_e b_e * q_e - sum_e b_e)
void calc_superblock(const uint a_offset, const uint b_offset, const uint v_im, const uint q_offset, const uint y_offset, const uint i, const uint num_blocks_per_row, const uint first_row, const uint num_rows) {
const uint y_idx = i * QUANT_K + y_offset;
[[unroll]] for (uint n = 0; n < num_rows; ++n) {
const uint ib0 = a_offset + (first_row+n)*num_blocks_per_row;
if (i >= num_blocks_per_row) {
continue;
}
const uint32_t qs_u32 = uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 8]) << 16);
const vec4 qs_u32_0 = vec4(unpack8(qs_u32 & 0x03030303));
const vec4 qs_u32_2 = vec4(unpack8((qs_u32 >> 2) & 0x03030303));
const vec4 qs_u32_4 = vec4(unpack8((qs_u32 >> 4) & 0x03030303));
const vec4 qs_u32_6 = vec4(unpack8((qs_u32 >> 6) & 0x03030303));
const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d);
[[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
vec2 b0 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 0]);
vec2 b16 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 8]);
vec2 b32 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 16]);
vec2 b48 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 24]);
vec2 b64 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 32]);
vec2 b80 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 40]);
vec2 b96 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 48]);
vec2 b112 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 56]);
FLOAT_TYPE sumq = FLOAT_TYPE(0.0);
FLOAT_TYPE sumb = FLOAT_TYPE(0.0);
[[unroll]] for (int l = 0; l < 2; ++l) {
sumq = fma(FLOAT_TYPE(b0[l]), FLOAT_TYPE(qs_u32_0[l ]),
fma(FLOAT_TYPE(b16[l]), FLOAT_TYPE(qs_u32_0[l+2]),
fma(FLOAT_TYPE(b32[l]), FLOAT_TYPE(qs_u32_2[l ]),
fma(FLOAT_TYPE(b48[l]), FLOAT_TYPE(qs_u32_2[l+2]),
fma(FLOAT_TYPE(b64[l]), FLOAT_TYPE(qs_u32_4[l ]),
fma(FLOAT_TYPE(b80[l]), FLOAT_TYPE(qs_u32_4[l+2]),
fma(FLOAT_TYPE(b96[l]), FLOAT_TYPE(qs_u32_6[l ]),
fma(FLOAT_TYPE(b112[l]), FLOAT_TYPE(qs_u32_6[l+2]), sumq))))))));
sumb += FLOAT_TYPE(b0[l]) + FLOAT_TYPE(b16[l]) + FLOAT_TYPE(b32[l]) + FLOAT_TYPE(b48[l])
+ FLOAT_TYPE(b64[l]) + FLOAT_TYPE(b80[l]) + FLOAT_TYPE(b96[l]) + FLOAT_TYPE(b112[l]);
}
temp[j][n] = fma(d, sumq - sumb, temp[j][n]);
}
}
}
void compute_outputs(const uint32_t first_row, const uint32_t num_rows) {
uint a_offset, b_offset, d_offset;
get_offsets(a_offset, b_offset, d_offset);
const uint num_blocks_per_row = p.ncols / QUANT_K;
// 16 threads are used to process each block
const uint it_size = gl_WorkGroupSize.x/16;
const uint tid = gl_LocalInvocationID.x;
const uint itid = tid%16; // 0...15
const uint ix = tid/16;
const uint v_im = itid/8; // 0 or 1. 0 computes 0..., 1 computes 128...
const uint v_in = itid - 8*v_im; // 0...7
const uint l0 = 2*v_in; // 0...15
const uint q_offset = 32*v_im + l0;
const uint y_offset = 128*v_im + l0;
[[unroll]] for (uint j = 0; j < NUM_COLS; ++j) {
[[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) {
temp[j][i] = FLOAT_TYPE(0);
}
}
for (uint i0 = 0; i0 < num_blocks_per_row; i0 += it_size)
calc_superblock(a_offset, b_offset, v_im, q_offset, y_offset, i0 + ix, num_blocks_per_row, first_row, num_rows);
reduce_result(temp, d_offset, first_row, num_rows, tid);
}
void main() {
const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z);
// do NUM_ROWS at a time, unless there aren't enough remaining rows
if (first_row + NUM_ROWS <= p.stride_d) {
compute_outputs(first_row, NUM_ROWS);
} else {
if (first_row >= p.stride_d) {
return;
}
compute_outputs(first_row, p.stride_d - first_row);
}
}

16
ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl
View file

@ -182,6 +182,22 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy);
buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw);
#elif defined(DATA_A_TQ2_0)
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
const uint ib = idx / 128; // 2 values per idx
const uint iqs = (idx % 128) * 2; // elem 0,2,4..254
const uint qsi = (iqs / 128) * 32 + (iqs % 32); // byte pair start
const uint shift = 2 * ((iqs % 128) / 32); // 0,2,4,6
const uvec2 qs = uvec2(data_a[ib].qs[qsi], data_a[ib].qs[qsi + 1]);
const float d = float(data_a[ib].d);
const vec2 v = d * (vec2((qs >> shift) & 3) - 1.0);
buf_a[buf_idx] = FLOAT_TYPEV2(v.xy);
#elif defined(DATA_A_Q3_K)
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;

24
ggml/src/ggml-vulkan/vulkan-shaders/types.glsl
View file

@ -303,6 +303,30 @@ struct block_q2_K_packed32
#define DATA_A_QUANT_K
#endif
#define QUANT_K_TQ2_0 256
// ternary (BitNet): 2-bit codes, w = (q - 1) * d; qs layout matches q2_K's
// two 32-byte groups with four bit-levels per byte
struct block_tq2_0
{
uint8_t qs[QUANT_K_TQ2_0/4];
float16_t d;
};
struct block_tq2_0_packed16
{
uint16_t qs[QUANT_K_TQ2_0/4/2];
float16_t d;
};
#if defined(DATA_A_TQ2_0)
#define QUANT_K QUANT_K_TQ2_0
#define QUANT_R 1
#define A_TYPE block_tq2_0
#define A_TYPE_PACKED16 block_tq2_0_packed16
#define DATA_A_QUANT_K
#endif
#define QUANT_K_Q3_K 256
struct block_q3_K

3
ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
View file

@ -86,6 +86,7 @@ const std::vector<std::string> type_names = {
"iq4_nl",
"mxfp4",
"nvfp4",
"tq2_0",
"bf16",
};
@ -759,7 +760,7 @@ void process_shaders() {
for (const auto& tname : type_names) {
// mul mat vec
std::string data_a_key = "DATA_A_" + to_uppercase(tname);
std::string shader = (string_ends_with(tname, "_k") || string_starts_with(tname, "iq1_") || string_starts_with(tname, "iq2_") || string_starts_with(tname, "iq3_")) ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp";
std::string shader = (string_ends_with(tname, "_k") || string_starts_with(tname, "iq1_") || string_starts_with(tname, "iq2_") || string_starts_with(tname, "iq3_") || tname == "tq2_0") ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp";
string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}));
string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}));

113
gguf-py/gguf/constants.py
View file

@ -407,6 +407,8 @@ class Keys:
class ClipGenAudio:
PROJECTOR_TYPE = "clip.gen.audio.projector_type" # for mixed modality models
# name of the weight variant, for settings that are not in the checkpoint
MODEL_VARIANT = "clip.gen.audio.model_variant"
EMBEDDING_LENGTH = "clip.gen.audio.embedding_length"
FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length"
BLOCK_COUNT = "clip.gen.audio.block_count"
@ -581,6 +583,7 @@ class MODEL_ARCH(IntEnum):
MELLUM = auto()
NANBEIGE = auto()
QWEN3TTS = auto()
POCKETTTS = auto()
class VISION_PROJECTOR_TYPE(IntEnum):
@ -1040,6 +1043,38 @@ class MODEL_TENSOR(IntEnum):
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
# pocket-tts: SEANet encoder (speaker path) and decoder (a.gen.wav path)
A_ENC_SEANET_CONV_IN = auto()
A_ENC_SEANET_CONV_OUT = auto()
A_ENC_SEANET_RES_CONV1 = auto() # residual unit, dilated conv
A_ENC_SEANET_RES_CONV2 = auto() # residual unit, pointwise conv
A_ENC_SEANET_SCALE_CONV = auto() # strided downsample conv
A_ENC_ATTN_SCALE = auto() # layer scale (gamma) on the attn output
A_ENC_FFN_SCALE_LS = auto() # layer scale (gamma) on the FFN output
A_ENC_SPEAKER_PROJ = auto() # voice latent -> backbone embd
A_GEN_FLOW_INPUT_PROJ = auto()
A_GEN_FLOW_COND_EMBD = auto()
A_GEN_FLOW_TIME_FREQS = auto() # timestep embedder, stored cos/sin frequencies
A_GEN_FLOW_TIME_UP = auto()
A_GEN_FLOW_TIME_DOWN = auto()
A_GEN_FLOW_TIME_NORM = auto() # RMSNorm alpha
A_GEN_FLOW_BLK_NORM = auto() # AdaLN res block, in_ln
A_GEN_FLOW_BLK_UP = auto()
A_GEN_FLOW_BLK_DOWN = auto()
A_GEN_FLOW_BLK_ADA = auto() # AdaLN modulation, -> shift/scale/gate
A_GEN_FLOW_FINAL_ADA = auto() # final layer AdaLN modulation, -> shift/scale
A_GEN_FLOW_FINAL_PROJ = auto()
A_GEN_OUT_EOS = auto() # end-of-speech head on the backbone hidden state
A_GEN_INPUT_LINEAR = auto() # generated latent -> backbone embd
A_GEN_EMB_MEAN = auto() # latent denormalization stats
A_GEN_EMB_STD = auto()
A_GEN_WAV_QUANT_OUT = auto() # DummyQuantizer output_proj, latent -> decoder dim
A_GEN_WAV_UPSAMPLE = auto() # frame rate -> encoder frame rate, depthwise convtr
A_GEN_WAV_SEANET_CONV_IN = auto()
A_GEN_WAV_SEANET_CONV_OUT = auto() # -> 1-channel PCM
A_GEN_WAV_SEANET_RES_CONV1 = auto()
A_GEN_WAV_SEANET_RES_CONV2 = auto()
A_GEN_WAV_SEANET_SCALE_CONV = auto() # strided upsample convtr
A_MMPROJ = auto()
A_MMPROJ_FC = auto()
A_MM_NORM_PRE = auto()
@ -1255,6 +1290,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
MODEL_ARCH.MELLUM: "mellum",
MODEL_ARCH.NANBEIGE: "nanbeige",
MODEL_ARCH.QWEN3TTS: "qwen3tts",
MODEL_ARCH.POCKETTTS: "pockettts",
}
VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = {
@ -1709,6 +1745,37 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
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_ENC_SEANET_CONV_IN: "a.seanet.conv_in",
MODEL_TENSOR.A_ENC_SEANET_CONV_OUT: "a.seanet.conv_out",
MODEL_TENSOR.A_ENC_SEANET_RES_CONV1: "a.seanet.blk.{bid}.res_conv1",
MODEL_TENSOR.A_ENC_SEANET_RES_CONV2: "a.seanet.blk.{bid}.res_conv2",
MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV: "a.seanet.blk.{bid}.scale_conv",
MODEL_TENSOR.A_ENC_ATTN_SCALE: "a.blk.{bid}.ls1",
MODEL_TENSOR.A_ENC_FFN_SCALE_LS: "a.blk.{bid}.ls2",
MODEL_TENSOR.A_ENC_SPEAKER_PROJ: "a.speaker_proj",
MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ: "a.gen.flow.input_proj",
MODEL_TENSOR.A_GEN_FLOW_COND_EMBD: "a.gen.flow.cond_embd",
MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS: "a.gen.flow.time.{bid}.freqs",
MODEL_TENSOR.A_GEN_FLOW_TIME_UP: "a.gen.flow.time.{bid}.up",
MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN: "a.gen.flow.time.{bid}.down",
MODEL_TENSOR.A_GEN_FLOW_TIME_NORM: "a.gen.flow.time.{bid}.norm",
MODEL_TENSOR.A_GEN_FLOW_BLK_NORM: "a.gen.flow.blk.{bid}.norm",
MODEL_TENSOR.A_GEN_FLOW_BLK_UP: "a.gen.flow.blk.{bid}.up",
MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN: "a.gen.flow.blk.{bid}.down",
MODEL_TENSOR.A_GEN_FLOW_BLK_ADA: "a.gen.flow.blk.{bid}.ada",
MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA: "a.gen.flow.final.ada",
MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ: "a.gen.flow.final.proj",
MODEL_TENSOR.A_GEN_OUT_EOS: "a.gen.out_eos",
MODEL_TENSOR.A_GEN_INPUT_LINEAR: "a.gen.input_linear",
MODEL_TENSOR.A_GEN_EMB_MEAN: "a.gen.emb_mean",
MODEL_TENSOR.A_GEN_EMB_STD: "a.gen.emb_std",
MODEL_TENSOR.A_GEN_WAV_QUANT_OUT: "a.gen.wav.quant_out",
MODEL_TENSOR.A_GEN_WAV_UPSAMPLE: "a.gen.wav.upsample",
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN: "a.gen.wav.seanet.conv_in",
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT: "a.gen.wav.seanet.conv_out",
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1: "a.gen.wav.seanet.blk.{bid}.res_conv1",
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2: "a.gen.wav.seanet.blk.{bid}.res_conv2",
MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV: "a.gen.wav.seanet.blk.{bid}.scale_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",
@ -2020,6 +2087,37 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
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_SEANET_CONV_IN,
MODEL_TENSOR.A_ENC_SEANET_CONV_OUT,
MODEL_TENSOR.A_ENC_SEANET_RES_CONV1,
MODEL_TENSOR.A_ENC_SEANET_RES_CONV2,
MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV,
MODEL_TENSOR.A_ENC_ATTN_SCALE,
MODEL_TENSOR.A_ENC_FFN_SCALE_LS,
MODEL_TENSOR.A_ENC_SPEAKER_PROJ,
MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ,
MODEL_TENSOR.A_GEN_FLOW_COND_EMBD,
MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS,
MODEL_TENSOR.A_GEN_FLOW_TIME_UP,
MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN,
MODEL_TENSOR.A_GEN_FLOW_TIME_NORM,
MODEL_TENSOR.A_GEN_FLOW_BLK_NORM,
MODEL_TENSOR.A_GEN_FLOW_BLK_UP,
MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN,
MODEL_TENSOR.A_GEN_FLOW_BLK_ADA,
MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA,
MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ,
MODEL_TENSOR.A_GEN_OUT_EOS,
MODEL_TENSOR.A_GEN_INPUT_LINEAR,
MODEL_TENSOR.A_GEN_EMB_MEAN,
MODEL_TENSOR.A_GEN_EMB_STD,
MODEL_TENSOR.A_GEN_WAV_QUANT_OUT,
MODEL_TENSOR.A_GEN_WAV_UPSAMPLE,
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN,
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT,
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1,
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2,
MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV,
MODEL_TENSOR.A_ENC_CONV_NORM_MEAN,
MODEL_TENSOR.A_ENC_CONV_NORM_VAR,
MODEL_TENSOR.A_ENC_MEL_FILTERS,
@ -4628,6 +4726,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.D2T,
],
MODEL_ARCH.DFLASH: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
MODEL_TENSOR.ATTN_NORM,
MODEL_TENSOR.ATTN_Q,
@ -4903,6 +5002,18 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
],
MODEL_ARCH.POCKETTTS: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
MODEL_TENSOR.ATTN_NORM,
MODEL_TENSOR.ATTN_Q,
MODEL_TENSOR.ATTN_K,
MODEL_TENSOR.ATTN_V,
MODEL_TENSOR.ATTN_OUT,
MODEL_TENSOR.FFN_NORM,
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
],
}
# tensors that will not be serialized
@ -5179,6 +5290,8 @@ class VisionProjectorType:
NEMOTRON_V2_VL = "nemotron_v2_vl"
QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder
QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor
POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder
POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder
HUNYUANVL = "hunyuanvl"
PARAKEET = "parakeet" # audio
MINIMAXM3 = "minimax_m3"

3
gguf-py/gguf/gguf_writer.py
View file

@ -1453,6 +1453,9 @@ class GGUFWriter:
def add_gen_audio_attention_layernorm_eps(self, value: float) -> None:
self.add_float32(Keys.ClipGenAudio.Attention.LAYERNORM_EPS, value)
def add_gen_audio_model_variant(self, value: str) -> None:
self.add_string(Keys.ClipGenAudio.MODEL_VARIANT, value)
def add_xielu_alpha_p(self, values: Sequence[float]):
self.add_array(Keys.xIELU.ALPHA_P, values)

211
models/templates/muse-glimmer.jinja Normal file
View file

@ -0,0 +1,211 @@
{#
Template: Muse Glimmer ATEM Chat Template
Renders the ATEM tool-calling protocol: reasoning channel (to=self), tool
channels (to=<tool>), and the user channel, plus tool definitions and the
valid-recipient list in the system block.
Whitespace note: every tag uses the {%- -%} / {{- -}} stripping markers, so
the indentation below is purely for readability and contributes nothing to
the rendered output.
#}
{%- macro render_content(content) -%}
{%- if content is string -%}
{{- content -}}
{%- elif content is not none -%}
{%- for part in content -%}
{%- if part['type'] == 'image' -%}
{{- '<|patch|>' -}}
{%- elif part['type'] == 'video' -%}
{{- '<|video|>' -}}
{%- elif part['type'] == 'text' -%}
{{- part['text'] -}}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
{%- endmacro -%}
{%- macro render_atem(tc) -%}
{%- set args = tc.function.arguments -%}
{%- if args is not mapping -%}
{{- raise_exception('Muse Glimmer ATEM chat template requires tool_call.function.arguments to be a dict (mapping); a JSON string cannot be parsed in the HF jinja sandbox.') -}}
{%- endif -%}
{{- '<atem:function_calls>\n<atem:invoke name="' + tc.function.name + '">\n' -}}
{%- for k, v in args.items() -%}
{{- '<atem:parameter name="' + k + '">' -}}
{%- if v is boolean -%}
{%- if v -%}
true
{%- else -%}
false
{%- endif -%}
{%- elif v is none -%}
null
{%- elif v is mapping or (v is iterable and v is not string) -%}
{{- v | tojson -}}
{%- else -%}
{{- v -}}
{%- endif -%}
{{- '</atem:parameter>\n' -}}
{%- endfor -%}
{{- '</atem:invoke>\n</atem:function_calls>' -}}
{%- endmacro -%}
{%- macro render_tool_defs(tools) -%}
{{- 'In this environment you have access to a set of tools you can use to answer the user\'s question.\n\n' -}}
{{- 'You can invoke a function by writing a "<atem:function_calls>" block like the following:\n' -}}
{{- '<atem:function_calls>\n<atem:invoke name="$FUNCTION_NAME">\n<atem:parameter name="$PARAMETER_NAME">$PARAMETER_VALUE</atem:parameter>\n...\n</atem:invoke>\n</atem:function_calls>\n\n' -}}
{{- 'String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions.\n' -}}
{{- 'Here are the functions available in JSONSchema format:\n' -}}
{{- '// Tool metadata\n' -}}
{%- set nsns = namespace(seen=[]) -%}
{%- for tool in tools -%}
{%- set fn = tool.function if tool.function is defined else tool -%}
{%- set tns = fn.name.split('.')[0] -%}
{%- if tns not in nsns.seen -%}
{%- set nsns.seen = nsns.seen + [tns] -%}
{%- endif -%}
{%- endfor -%}
{%- set nd = tool_namespace_descriptions if tool_namespace_descriptions is defined else {} -%}
{%- for tns in nsns.seen -%}
{{- '{"name": ' + (tns | tojson) + ', "description": ' + ((nd[tns] if tns in nd else '') | tojson) + '}\n' -}}
{%- endfor -%}
{{- '// Function schemas' -}}
{%- for tool in tools -%}
{%- set fn = tool.function if tool.function is defined else tool -%}
{{- '\n{"name": ' + (fn.name | tojson) + ', "description": ' + (fn.description | tojson) + ', "parameters": ' + (fn.parameters | tojson) + '}' -}}
{%- endfor -%}
{{- '\n\nHere\'s an example of how to call a function in the tool set:\n' -}}
{{- '(If the tool namespace is not specified, invoke the function directly as `example_function_name` rather than `example_tool_name.example_function_name`)\n\n' -}}
{{- 'to=example_tool_name.example_function_name\n\n' -}}
{{- '<atem:function_calls>\n<atem:invoke name="example_tool_name.example_function_name">\n' -}}
{{- '<atem:parameter name="example_parameter_1">value_1</atem:parameter>\n' -}}
{{- '<atem:parameter name="example_parameter_2">This is the value for the second parameter\nthat can span\n"multiple" lines\n</atem:parameter>\n' -}}
{{- '</atem:invoke>\n</atem:function_calls>' -}}
{%- endmacro -%}
{%- macro render_reasoning() -%}
{%- set rs = reasoning_strength if reasoning_strength is defined and reasoning_strength else 'high' -%}
{{- 'Reasoning strength: ' + rs + '.' -}}
{%- endmacro -%}
{%- macro render_system_meta(tools) -%}
{%- set rns = namespace(recipients=['"self"'], nslist=[]) -%}
{%- if tools -%}
{%- for tool in tools -%}
{%- set fn = tool.function if tool.function is defined else tool -%}
{%- set tns = fn.name.split('.')[0] -%}
{%- if tns not in rns.nslist -%}
{%- set rns.nslist = rns.nslist + [tns] -%}
{%- endif -%}
{%- endfor -%}
{%- for tns in rns.nslist -%}
{%- set rns.recipients = rns.recipients + ['"' + tns + '.*"'] -%}
{%- endfor -%}
{%- endif -%}
{%- set rns.recipients = rns.recipients + ['"user"'] -%}
{{- '# Valid recipients: ' + rns.recipients | join(', ') + '.' -}}
{%- endmacro -%}
{{- bos_token -}}
{%- set ns = namespace(has_system=false) -%}
{%- for m in messages -%}
{%- if m['role'] == 'system' -%}
{%- set ns.has_system = true -%}
{%- endif -%}
{%- endfor -%}
{%- if not ns.has_system -%}
{{- '<|start|>system<|message|>You are a helpful AI assistant.' -}}
{%- set kc = knowledge_cutoff if knowledge_cutoff is defined and knowledge_cutoff else '2026-01-04' -%}
{{- '\nKnowledge cutoff: ' + kc + '.' -}}
{%- if current_date is defined and current_date -%}
{{- '\nCurrent date: ' + current_date + '.' -}}
{%- elif strftime_now is defined -%}
{{- '\nCurrent date: ' + strftime_now('%Y-%m-%d') + '.' -}}
{%- endif -%}
{{- '\n\n' -}}
{{- render_reasoning() -}}
{%- if tools -%}
{{- '\n\n' -}}
{{- render_tool_defs(tools) -}}
{%- endif -%}
{{- '\n\n' -}}
{{- render_system_meta(tools) -}}
{{- '<|eot|>' -}}
{%- endif -%}
{%- for message in messages -%}
{%- set role = message['role'] -%}
{%- set end_token = '<|eom|>' if (not loop.last and messages[loop.index0 + 1]['role'] == role) else '<|eot|>' -%}
{%- if role == 'system' -%}
{#- Callers sometimes write the directive into the system prompt themselves.
Normalise "Reasoning effort" to "Reasoning strength" (jinja has no
case-insensitive replace, hence the four realistic casings), then skip
the kwarg-driven line below if the prompt already carries one. -#}
{%- set sys_text = render_content(message['content'])
| replace('Reasoning effort', 'Reasoning strength')
| replace('Reasoning Effort', 'Reasoning Strength')
| replace('reasoning effort', 'reasoning strength')
| replace('REASONING EFFORT', 'REASONING STRENGTH') -%}
{{- '<|start|>system<|message|>' -}}
{{- sys_text -}}
{%- if 'reasoning strength' not in (sys_text | lower) -%}
{{- '\n\n' -}}
{{- render_reasoning() -}}
{%- endif -%}
{%- if tools -%}
{{- '\n\n' -}}
{{- render_tool_defs(tools) -}}
{%- endif -%}
{{- '\n\n' -}}
{{- render_system_meta(tools) -}}
{{- '<|eot|>' -}}
{%- elif role == 'user' -%}
{{- '<|start|>user<|message|>' -}}
{{- render_content(message['content']) -}}
{{- '<|eot|>' -}}
{%- elif role == 'tool' -%}
{%- set tname = message.get('name') -%}
{%- if not tname -%}
{%- set tcid = message.get('tool_call_id') -%}
{%- set rns = namespace(name=tcid if tcid else '') -%}
{%- for m in messages -%}
{%- if m.get('tool_calls') -%}
{%- for tc in m['tool_calls'] -%}
{%- if tcid is not none and tc.id is defined and tc.id == tcid -%}
{%- set rns.name = tc.function.name -%}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
{%- endfor -%}
{%- set tname = rns.name -%}
{%- endif -%}
{{- '<|start|>tool ' + tname + '<|message|><tool_output name="' + tname + '">\n' -}}
{{- render_content(message['content']) -}}
{{- '\n</tool_output><|eot|>' -}}
{%- elif role == 'assistant' -%}
{%- if message.get('reasoning_content') -%}
{{- '<|start|>assistant to=self<|message|>' + message['reasoning_content'] + '<|eom|>' -}}
{%- endif -%}
{%- if message.get('tool_calls') -%}
{%- for tc in message['tool_calls'] -%}
{{- '<|start|>assistant to=' + tc.function.name + '<|message|>' -}}
{{- render_atem(tc) -}}
{%- if loop.last -%}
{{- end_token -}}
{%- else -%}
{{- '<|eom|>' -}}
{%- endif -%}
{%- endfor -%}
{%- else -%}
{%- set recipient = message.get('recipient') or 'user' -%}
{%- set end_turn = message.get('end_turn') -%}
{%- if end_turn is none -%}
{%- set end_turn = not (recipient and recipient != 'user') -%}
{%- endif -%}
{{- '<|start|>assistant' -}}
{%- if recipient -%}
{{- ' to=' + recipient -}}
{%- endif -%}
{{- '<|message|>' -}}
{{- render_content(message['content']) -}}
{{- ('<|eot|>' if end_turn else '<|eom|>') -}}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
{{- '<|start|>assistant' -}}
{%- endif -%}

1
src/llama-arch.cpp
View file

@ -147,6 +147,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_NANBEIGE, "nanbeige" },
{ LLM_ARCH_QWEN3TTS, "qwen3tts" },
{ LLM_ARCH_POCKETTTS, "pockettts" },
{ LLM_ARCH_UNKNOWN, "(unknown)" },
};

1
src/llama-arch.h
View file

@ -152,6 +152,7 @@ enum llm_arch {
LLM_ARCH_DFLASH,
LLM_ARCH_NANBEIGE,
LLM_ARCH_QWEN3TTS,
LLM_ARCH_POCKETTTS,
LLM_ARCH_UNKNOWN,
};

7
src/llama-model.cpp
View file

@ -151,6 +151,7 @@
#include "models/plamo2.cpp"
#include "models/plamo3.cpp"
#include "models/plm.cpp"
#include "models/pockettts.cpp"
#include "models/qwen.cpp"
#include "models/qwen2.cpp"
#include "models/qwen2moe.cpp"
@ -261,6 +262,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_qwen3vlmoe(params);
case LLM_ARCH_QWEN3TTS:
return new llama_model_qwen3tts(params);
case LLM_ARCH_POCKETTTS:
return new llama_model_pockettts(params);
case LLM_ARCH_PHI2:
return new llama_model_phi2(params);
case LLM_ARCH_PHI3:
@ -1265,6 +1268,9 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_CONVNEXT_EMBEDDING_LENGTH, hparams.convnext.n_embd);
ml.get_key(LLM_KV_CONVNEXT_BLOCK_COUNT, hparams.convnext.n_layer);
GGML_ASSERT(hparams.posnet.n_layer <= hparams.n_layer_all);
GGML_ASSERT(hparams.convnext.n_layer <= hparams.n_layer_all);
}
GGML_ASSERT(hparams.n_expert <= LLAMA_MAX_EXPERTS);
@ -2782,6 +2788,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_MAINCODER:
case LLM_ARCH_GLM_DSA:
case LLM_ARCH_NANBEIGE:
case LLM_ARCH_POCKETTTS:
return LLAMA_ROPE_TYPE_NORM;
// the pairs of head values are offset by n_rot/2

5
src/llama-model.h
View file

@ -623,8 +623,9 @@ struct llama_model {
struct ggml_tensor * per_layer_model_proj = nullptr;
struct ggml_tensor * per_layer_proj_norm = nullptr;
// eagle3
struct ggml_tensor * fc = nullptr; // feature fusion layer
// eagle3 / dflash feature fusion layer
struct ggml_tensor * fc = nullptr;
struct ggml_tensor * fc_s = nullptr;
struct ggml_tensor * d2t = nullptr; // draft to target vocabulary mapping
// dspark

26
src/models/dflash.cpp
View file

@ -79,6 +79,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
// DSpark = DFlash + a semi-autoregressive Markov head and Confidence head
//
// TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4)
@ -97,6 +98,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
}
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);
fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm
@ -205,7 +207,7 @@ template <>
llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
ggml_tensor * cur = build_inp_embd_enc();
cur = build_lora_mm(model.fc, cur);
cur = build_lora_mm(model.fc, cur, model.fc_s);
cb(cur, "fc_out", -1);
cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
@ -460,9 +462,9 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
cb(cur, "ffn_norm", il);
cur = build_ffn(cur,
layer.ffn_up, NULL, NULL,
layer.ffn_gate, NULL, NULL,
layer.ffn_down, NULL, NULL,
layer.ffn_up, NULL, layer.ffn_up_s,
layer.ffn_gate, NULL, layer.ffn_gate_s,
layer.ffn_down, NULL, layer.ffn_down_s,
NULL,
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(cur, "ffn_out", il);
@ -479,15 +481,17 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
res->t_embd = cur;
// lm_head from the target model (shared via ctx_other)
auto * output = model.output;
auto * output = model.output;
auto * output_s = model.output_s;
if (output == nullptr) {
GGML_ASSERT(cparams.ctx_other != nullptr);
const auto * model_other = llama_get_model(cparams.ctx_other);
GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection");
output = model_other->output;
output = model_other->output;
output_s = model_other->output_s;
}
cur = build_lora_mm(output, cur);
cur = build_lora_mm(output, cur, output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
@ -655,15 +659,17 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
cb(cur, "result_norm", -1);
// lm_head from the target model (shared via ctx_other)
auto * output = model.output;
auto * output = model.output;
auto * output_s = model.output_s;
if (output == nullptr) {
GGML_ASSERT(cparams.ctx_other != nullptr);
const auto * model_other = llama_get_model(cparams.ctx_other);
GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
output = model_other->output;
output = model_other->output;
output_s = model_other->output_s;
}
cur = build_lora_mm(output, cur);
cur = build_lora_mm(output, cur, output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;

13
src/models/models.h
View file

@ -713,6 +713,19 @@ struct llama_model_gpt2 : public llama_model_base {
};
struct llama_model_pockettts : public llama_model_base {
llama_model_pockettts(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
graph(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_codeshell : public llama_model_base {
llama_model_codeshell(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;

12
src/models/nemotron-h.cpp
View file

@ -177,8 +177,11 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
auto * inp = build_inp_mem_hybrid();
ggml_tensor * inp_out_ids = build_inp_out_ids();
const bool extract_final_inp = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer];
for (int il = 0; il < n_layer; ++il) {
res->t_layer_inp[il] = inpL;
struct ggml_tensor * inpSA = inpL;
// norm
@ -195,7 +198,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
cur = build_ffn_layer(cur, model, il);
}
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked && !extract_final_inp) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -209,6 +212,13 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
}
cur = inpL;
if (extract_final_inp) {
res->t_layer_inp[n_layer] = cur;
if (inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
}
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);

146
src/models/pockettts.cpp Normal file
View file

@ -0,0 +1,146 @@
#include "models.h"
// backbone of the pocket-tts CALM pipeline: the "text" side of a flow language model.
// it has no lm_head, the audio latents are produced by the flow net inside the mmproj
void llama_model_pockettts::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
switch (hparams.n_layer()) {
case 6: type = LLM_TYPE_109M; break;
case 24: type = LLM_TYPE_335M; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_pockettts::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"), {n_embd}, 0);
// no output head, the logits are unused; reuse the embedding table so a sampler can still run
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0);
create_tensor_qkv(layer, i, n_embd, n_embd, n_embd_gqa, n_embd_gqa, TENSOR_NOT_REQUIRED);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", i), {n_embd}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
}
}
std::unique_ptr<llm_graph_context> llama_model_pockettts::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
llama_model_pockettts::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
GGML_ASSERT(n_embd_head == n_rot);
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
ggml_tensor * inp_pos = build_inp_pos();
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * inp_out_ids = build_inp_out_ids();
for (int il = 0; il < n_layer; ++il) {
cur = build_norm(inpL,
model.layers[il].attn_norm,
model.layers[il].attn_norm_b,
LLM_NORM, il);
cb(cur, "attn_norm", il);
// self-attention
{
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
n_embd_head, n_head, n_head_kv, il);
Qcur = ggml_rope_ext(
ctx0, Qcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow
);
Kcur = ggml_rope_ext(
ctx0, Kcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow
);
cb(Qcur, "Qcur", il);
cb(Kcur, "Kcur", il);
cb(Vcur, "Vcur", il);
cur = build_attn(inp_attn,
model.layers[il].wo, NULL, model.layers[il].wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
}
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
cb(ffn_inp, "ffn_inp", il);
// FF
{
cur = build_norm(ffn_inp,
model.layers[il].ffn_norm,
model.layers[il].ffn_norm_b,
LLM_NORM, il);
cb(cur, "ffn_norm", il);
cur = build_ffn(cur,
model.layers[il].ffn_up, NULL, NULL,
NULL, NULL, NULL,
model.layers[il].ffn_down, NULL, NULL,
NULL,
LLM_FFN_GELU, LLM_FFN_SEQ, il);
cb(cur, "ffn_out", il);
}
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
// input for next layer
inpL = cur;
}
cur = build_norm(inpL,
model.output_norm,
model.output_norm_b,
LLM_NORM, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = build_lora_mm(model.output, cur, model.output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}

15
tools/mtmd/README-dev.md
View file

@ -59,8 +59,10 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i
### 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:
1. Make sure to consult merged PRs about adding new TTS models, especially reviewer comments
- Example: https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr+mtmd+tts+is%3Amerged
2. Establish a list of reusable and missing components from the current mtmd implementation.
3. 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`)
@ -70,12 +72,17 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i
- 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:
- For GGUF metadata:
- Reuse as many existing keys as possible
- In most cases, you can hard-code model configs in the model graph class, or in `clip_hparams`
- If some values need to be exposed to the `mtmd_helper` layer, hard-code them in `mtmd_helper` and distinguish by pipeline and `mtmd_gen_audio_info::model_variant` if necessary
- Do NOT add new GGUF metadata or new fields to `mtmd_gen_audio_info` unless you can prove that you absolutely need them
4. 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`
5. 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**.

40
tools/mtmd/clip-impl.h
View file

@ -92,7 +92,9 @@
#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"
// name of the weight variant, for settings that are not in the checkpoint
#define KEY_GEN_AUDIO_VARIANT "clip.gen.audio.model_variant"
#define KEY_AUDIO_SUBSMPL_FACTOR "clip.audio.subsampling_factor"
//
// tensor name constants
@ -246,6 +248,38 @@
#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"
// pocket-tts
#define TN_A_SEANET_CONV_IN "a.seanet.conv_in.%s"
#define TN_A_SEANET_CONV_OUT "a.seanet.conv_out.%s"
#define TN_A_SEANET_RES_CONV1 "a.seanet.blk.%d.res_conv1.%s"
#define TN_A_SEANET_RES_CONV2 "a.seanet.blk.%d.res_conv2.%s"
#define TN_A_SEANET_SCALE_CONV "a.seanet.blk.%d.scale_conv.%s"
#define TN_A_SPEAKER_PROJ "a.speaker_proj.%s"
#define TN_A_DOWNSAMPLE_CONV "a.downsample.conv.%s"
#define TN_A_GEN_FLOW_INPUT_PROJ "a.gen.flow.input_proj.%s"
#define TN_A_GEN_FLOW_COND_EMBD "a.gen.flow.cond_embd.%s"
#define TN_A_GEN_FLOW_TIME_FREQS "a.gen.flow.time.%d.freqs"
#define TN_A_GEN_FLOW_TIME_UP "a.gen.flow.time.%d.up.%s"
#define TN_A_GEN_FLOW_TIME_DOWN "a.gen.flow.time.%d.down.%s"
#define TN_A_GEN_FLOW_TIME_NORM "a.gen.flow.time.%d.norm"
#define TN_A_GEN_FLOW_BLK_NORM "a.gen.flow.blk.%d.norm.%s"
#define TN_A_GEN_FLOW_BLK_UP "a.gen.flow.blk.%d.up.%s"
#define TN_A_GEN_FLOW_BLK_DOWN "a.gen.flow.blk.%d.down.%s"
#define TN_A_GEN_FLOW_BLK_ADA "a.gen.flow.blk.%d.ada.%s"
#define TN_A_GEN_FLOW_FINAL_ADA "a.gen.flow.final.ada.%s"
#define TN_A_GEN_FLOW_FINAL_PROJ "a.gen.flow.final.proj.%s"
#define TN_A_GEN_OUT_EOS "a.gen.out_eos.%s"
#define TN_A_GEN_INPUT_LINEAR "a.gen.input_linear.%s"
#define TN_A_GEN_EMB_MEAN "a.gen.emb_mean"
#define TN_A_GEN_EMB_STD "a.gen.emb_std"
#define TN_A_GEN_WAV_QUANT_OUT "a.gen.wav.quant_out.%s"
#define TN_A_GEN_WAV_UPSAMPLE "a.gen.wav.upsample.%s"
#define TN_A_GEN_WAV_SEANET_CONV_IN "a.gen.wav.seanet.conv_in.%s"
#define TN_A_GEN_WAV_SEANET_CONV_OUT "a.gen.wav.seanet.conv_out.%s"
#define TN_A_GEN_WAV_SEANET_RES_CONV1 "a.gen.wav.seanet.blk.%d.res_conv1.%s"
#define TN_A_GEN_WAV_SEANET_RES_CONV2 "a.gen.wav.seanet.blk.%d.res_conv2.%s"
#define TN_A_GEN_WAV_SEANET_SCALE_CONV "a.gen.wav.seanet.blk.%d.scale_conv.%s"
// cogvlm
#define TN_MM_POST_FC_NORM "mm.post_fc_norm.%s"
#define TN_MM_H_TO_4H "mm.up.%s"
@ -455,6 +489,8 @@ enum projector_type {
PROJECTOR_TYPE_MIMO_AUDIO,
PROJECTOR_TYPE_QWEN3TTS_SPKENC,
PROJECTOR_TYPE_QWEN3TTS_GEN,
PROJECTOR_TYPE_POCKETTTS_SPKENC,
PROJECTOR_TYPE_POCKETTTS_GEN,
PROJECTOR_TYPE_MUSE_GLIMMER,
PROJECTOR_TYPE_UNKNOWN,
};
@ -515,6 +551,8 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
{ PROJECTOR_TYPE_PARAKEET, "parakeet"},
{ PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"},
{ PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"},
{ PROJECTOR_TYPE_POCKETTTS_SPKENC, "pockettts_spkenc"},
{ PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"},
{ PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"},
};

89
tools/mtmd/clip-model.h
View file

@ -141,6 +141,20 @@ 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)
// threshold for the "out_eos_score" graph output
float gen_eos_threshold = 0.0f;
// name of the weight variant, some pipelines tune themselves on it
std::string gen_model_variant;
// pocket-tts
static constexpr int32_t pockettts_max_spk_seconds = 30;
int32_t seanet_n_stage = 0;
std::vector<int32_t> seanet_ratios; // encoder order (reversed compared to the config)
int32_t mimi_downsample = 0; // encoder frame rate / model frame rate
int32_t mimi_tfm_context = 0; // attention window of the mimi transformers, in frames
int32_t flow_n_step = 1; // lsd_decode steps
// qwen3tts code2wav
int32_t wav_tfm_n_layer = 0;
int32_t wav_tfm_n_embd = 0;
@ -402,6 +416,63 @@ struct qf_block {
std::vector<clip_layer> qf_proj_layers;
};
// pocket-tts SEANet stack, used in both directions:
// encoder = conv_in -> per stage (residual unit, strided conv) -> conv_out
// decoder = conv_in -> per stage (strided convtr, residual unit) -> conv_out
struct clip_seanet {
// one residual unit: ELU -> dilated conv -> ELU -> pointwise conv, added to the input
struct stage {
ggml_tensor * res_conv1_w = nullptr;
ggml_tensor * res_conv1_b = nullptr;
ggml_tensor * res_conv2_w = nullptr;
ggml_tensor * res_conv2_b = nullptr;
ggml_tensor * scale_conv_w = nullptr; // strided conv (encoder) or convtr (decoder)
ggml_tensor * scale_conv_b = nullptr;
};
ggml_tensor * conv_in_w = nullptr;
ggml_tensor * conv_in_b = nullptr;
ggml_tensor * conv_out_w = nullptr;
ggml_tensor * conv_out_b = nullptr;
std::vector<stage> stages;
};
// pocket-tts flow-matching decoder (SimpleMLPAdaLN)
struct clip_flow_net {
// AdaLN res block: in_ln -> modulate -> Linear -> SiLU -> Linear, gated residual
struct block {
ggml_tensor * norm_w = nullptr;
ggml_tensor * norm_b = nullptr;
ggml_tensor * up_w = nullptr;
ggml_tensor * up_b = nullptr;
ggml_tensor * down_w = nullptr;
ggml_tensor * down_b = nullptr;
ggml_tensor * ada_w = nullptr; // -> shift, scale, gate
ggml_tensor * ada_b = nullptr;
};
// timestep embedder: cos/sin(t * freqs) -> Linear -> SiLU -> Linear -> RMSNorm
struct time_embd {
ggml_tensor * freqs = nullptr;
ggml_tensor * up_w = nullptr;
ggml_tensor * up_b = nullptr;
ggml_tensor * down_w = nullptr;
ggml_tensor * down_b = nullptr;
ggml_tensor * norm = nullptr; // RMSNorm alpha
};
ggml_tensor * input_proj_w = nullptr;
ggml_tensor * input_proj_b = nullptr;
ggml_tensor * cond_embd_w = nullptr;
ggml_tensor * cond_embd_b = nullptr;
ggml_tensor * final_ada_w = nullptr; // -> shift, scale
ggml_tensor * final_ada_b = nullptr;
ggml_tensor * final_proj_w = nullptr;
ggml_tensor * final_proj_b = nullptr;
std::vector<time_embd> time;
std::vector<block> blocks;
};
// qwen3tts code2wav: RVQ codes -> raw PCM
struct clip_code2wav {
// "upsample" stage: one ConvNeXt block plus the causal ConvTranspose1d before it
@ -699,6 +770,24 @@ struct clip_model {
// qwen3tts code2wav: RVQ codes -> raw PCM
clip_code2wav c2w;
// pocket-tts: SEANet stack, shared by the encoder (speaker path) and the decoder (gen path)
clip_seanet seanet;
// pocket-tts: voice latent -> backbone embd (speaker path)
ggml_tensor * spk_proj_w = nullptr;
ggml_tensor * downsample_w = nullptr;
// pocket-tts: flow-matching decoder, backbone hidden state -> next latent
clip_flow_net flow;
ggml_tensor * gen_out_eos_w = nullptr;
ggml_tensor * gen_out_eos_b = nullptr;
ggml_tensor * gen_input_lin_w = nullptr; // latent -> backbone embd
ggml_tensor * gen_emb_mean = nullptr;
ggml_tensor * gen_emb_std = nullptr;
ggml_tensor * gen_quant_out_w = nullptr; // latent -> decoder dim
ggml_tensor * gen_upsample_w = nullptr; // depthwise convtr, frame rate -> encoder frame rate
std::vector<clip_layer> gen_tfm_layers; // mimi decoder_transformer
// cogvlm
ggml_tensor * mm_post_fc_norm_w = nullptr;
ggml_tensor * mm_post_fc_norm_b = nullptr;

306
tools/mtmd/clip.cpp
View file

@ -64,6 +64,9 @@
#include "models/paddleocr.cpp"
#include "models/parakeet.cpp"
#include "models/pixtral.cpp"
#include "models/pockettts-gen.cpp"
#include "models/pockettts-seanet.cpp"
#include "models/pockettts-spkenc.cpp"
#include "models/qwen2vl.cpp"
#include "models/qwen3vl.cpp"
#include "models/qwen3a.cpp"
@ -228,6 +231,10 @@ struct clip_ctx {
bool support_batch = false;
// for audio gen, reseeded only when the caller asks for another seed
std::mt19937 rng{std::random_device{}()};
uint32_t rng_seed = UINT32_MAX;
clip_ctx(clip_context_params & ctx_params) {
flash_attn_type = ctx_params.flash_attn_type;
no_alloc = ctx_params.no_alloc;
@ -1113,6 +1120,25 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
{
builder = std::make_unique<clip_graph_qwen3tts_spkenc>(ctx, img);
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
{
builder = std::make_unique<clip_graph_pockettts_spkenc>(ctx, img);
} break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
{
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
const int n_step = ctx->model.hparams.flow_n_step;
const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0];
GGML_ASSERT(n_step > 0);
GGML_ASSERT(n_latent > 0);
// "inp_feats" takes the caller's buffer as-is, the graph must consume all of it
if (params && params->feats) {
GGML_ASSERT(params->feats->size() % (size_t) n_latent == 0);
GGML_ASSERT(params->feats->size() >= (size_t) n_latent);
}
const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1;
builder = std::make_unique<clip_graph_pockettts_gen>(ctx, img, gen_process, n_step, n_frames);
} break;
case PROJECTOR_TYPE_QWEN3TTS_GEN:
{
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
@ -1354,6 +1380,7 @@ struct clip_model_loader {
// these are unused, but still need to be set to avoid issues
hparams.image_size = 0;
hparams.patch_size = 1;
get_string(KEY_GEN_AUDIO_VARIANT, hparams.gen_model_variant, false);
} else {
GGML_ASSERT(false && "unknown modality");
@ -1498,7 +1525,7 @@ struct clip_model_loader {
} break;
case PROJECTOR_TYPE_PARAKEET:
{
get_u32(KEY_AUDIO_SUBSAMPLING_FACTOR, hparams.subsampling_factor);
get_u32(KEY_AUDIO_SUBSMPL_FACTOR, hparams.subsampling_factor);
GGML_ASSERT(hparams.subsampling_factor == 8 &&
"subsampling_factor must match the conv strides in clip_graph_parakeet::build()");
get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size);
@ -1827,6 +1854,22 @@ struct clip_model_loader {
// matches the reference decoder's sliding_window (speech_tokenizer/config.json)
hparams.wav_tfm_swa = 72;
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
case PROJECTOR_TYPE_POCKETTTS_GEN:
{
// mimi front-end takes the raw waveform, no mel
hparams.audio_sample_rate = 24000;
// seanet ratios are [6,5,4] in the config, the encoder reverses them
hparams.seanet_ratios = { 4, 5, 6 };
hparams.seanet_n_stage = (int32_t) hparams.seanet_ratios.size();
hparams.mimi_downsample = 16;
// matches the reference transformer's "context"
hparams.mimi_tfm_context = 250;
hparams.rope_theta = 10000.0f;
// flow_lm defaults, see pocket_tts/default_parameters.py
hparams.flow_n_step = 1;
hparams.gen_eos_threshold = -4.0f;
} break;
case PROJECTOR_TYPE_PADDLEOCR:
{
hparams.n_merge = 2;
@ -2029,7 +2072,9 @@ struct clip_model_loader {
// GEMMA4UA is encoder-free: it uses n_mel_bins as a raw-waveform frame size (640) and has no FFT/filterbank, so the mel-range and FFT
// checks below do not apply to it.
const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA;
// pocket-tts is encoder-free in the same sense: mimi convolves the raw waveform
const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA &&
model.proj_type != PROJECTOR_TYPE_POCKETTTS_SPKENC;
// Validate audio hparams loaded from GGUF metadata
if (hparams.n_mel_bins <= 0 || (fft_based && hparams.n_mel_bins > 256)) {
@ -2107,6 +2152,31 @@ struct clip_model_loader {
return cur;
};
// pocket-tts: the encoder and the decoder share the same layout, only the prefix differs
auto load_seanet = [&](clip_seanet & seanet, bool is_decoder) {
const char * conv_in = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_IN : TN_A_SEANET_CONV_IN;
const char * conv_out = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_OUT : TN_A_SEANET_CONV_OUT;
const char * res1 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV1 : TN_A_SEANET_RES_CONV1;
const char * res2 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV2 : TN_A_SEANET_RES_CONV2;
const char * scale = is_decoder ? TN_A_GEN_WAV_SEANET_SCALE_CONV : TN_A_SEANET_SCALE_CONV;
seanet.conv_in_w = get_tensor(string_format(conv_in, "weight"));
seanet.conv_in_b = get_tensor(string_format(conv_in, "bias"));
seanet.conv_out_w = get_tensor(string_format(conv_out, "weight"));
seanet.conv_out_b = get_tensor(string_format(conv_out, "bias"));
seanet.stages.resize(hparams.seanet_n_stage);
for (int i = 0; i < hparams.seanet_n_stage; i++) {
auto & stage = seanet.stages[i];
stage.res_conv1_w = get_tensor(string_format(res1, i, "weight"));
stage.res_conv1_b = get_tensor(string_format(res1, i, "bias"));
stage.res_conv2_w = get_tensor(string_format(res2, i, "weight"));
stage.res_conv2_b = get_tensor(string_format(res2, i, "bias"));
stage.scale_conv_w = get_tensor(string_format(scale, i, "weight"));
stage.scale_conv_b = get_tensor(string_format(scale, i, "bias"));
}
};
auto get_vector = [&](const std::string & name) {
std::vector<float> result;
auto it = tensor_offset.find(name);
@ -2168,7 +2238,8 @@ struct clip_model_loader {
const bool has_standard_layers = (
model.proj_type != PROJECTOR_TYPE_GEMMA3NV &&
model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC);
model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC &&
model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN);
// layers
const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0;
@ -2842,6 +2913,81 @@ struct clip_model_loader {
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_POCKETTTS_SPKENC:
{
load_seanet(model.seanet, false);
model.downsample_w = get_tensor(string_format(TN_A_DOWNSAMPLE_CONV, "weight"));
model.spk_proj_w = get_tensor(string_format(TN_A_SPEAKER_PROJ, "weight"));
} break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
{
auto & flow = model.flow;
flow.input_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "weight"));
flow.input_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "bias"));
flow.cond_embd_w = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "weight"));
flow.cond_embd_b = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "bias"));
flow.final_ada_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "weight"));
flow.final_ada_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "bias"));
flow.final_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "weight"));
flow.final_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "bias"));
flow.time.resize(2);
for (size_t i = 0; i < flow.time.size(); i++) {
auto & t = flow.time[i];
t.freqs = get_tensor(string_format(TN_A_GEN_FLOW_TIME_FREQS, (int) i));
t.up_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "weight"));
t.up_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "bias"));
t.down_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "weight"));
t.down_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "bias"));
t.norm = get_tensor(string_format(TN_A_GEN_FLOW_TIME_NORM, (int) i));
}
// one AdaLN block per flow depth, the count is only known from the tensors
for (int il = 0; ; il++) {
ggml_tensor * probe = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "weight"), false);
if (probe == nullptr) {
break;
}
clip_flow_net::block blk;
blk.norm_w = probe;
blk.norm_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "bias"));
blk.up_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "weight"));
blk.up_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "bias"));
blk.down_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "weight"));
blk.down_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "bias"));
blk.ada_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "weight"));
blk.ada_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "bias"));
flow.blocks.push_back(blk);
}
model.gen_out_eos_w = get_tensor(string_format(TN_A_GEN_OUT_EOS, "weight"));
model.gen_out_eos_b = get_tensor(string_format(TN_A_GEN_OUT_EOS, "bias"));
model.gen_input_lin_w = get_tensor(string_format(TN_A_GEN_INPUT_LINEAR, "weight"));
model.gen_emb_mean = get_tensor(TN_A_GEN_EMB_MEAN);
model.gen_emb_std = get_tensor(TN_A_GEN_EMB_STD);
// mimi decoder
model.gen_quant_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_OUT, "weight"));
model.gen_upsample_w = get_tensor(string_format(TN_A_GEN_WAV_UPSAMPLE, "weight"));
load_seanet(model.seanet, true);
model.gen_tfm_layers.resize(hparams.n_layer);
for (int il = 0; il < hparams.n_layer; il++) {
auto & layer = model.gen_tfm_layers[il];
const char * p = "a.gen.wav.tfm";
layer.ln_1_w = get_tensor(string_format(TN_LN_1, p, il, "weight"));
layer.ln_1_b = get_tensor(string_format(TN_LN_1, p, il, "bias"));
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.ls_1_w = get_tensor(string_format(TN_LS_1, p, il, "weight"));
layer.ln_2_w = get_tensor(string_format(TN_LN_2, p, il, "weight"));
layer.ln_2_b = get_tensor(string_format(TN_LN_2, p, il, "bias"));
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"));
layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight"));
}
} break;
case PROJECTOR_TYPE_QWEN3TTS_GEN:
{
// code_predictor
@ -4147,6 +4293,17 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
// one hidden-state vector fed back to the talker per call
n_patches = 1;
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
{
// one conditioning row per 12.5Hz frame
const int hop = ctx->model.hparams.mimi_downsample * 120;
n_patches = img->nx() / hop;
} break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
{
// one latent per call for GEN_CODE, GEN_WAV sizes its input from the caller
n_patches = 1;
} break;
case PROJECTOR_TYPE_GRANITE4_VISION:
{
// Per-tile output token count: each projector block outputs
@ -4188,6 +4345,15 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
return clip_encode(ctx, &params);
}
// persisted state slots of the gen-audio decoder, per pipeline
static std::vector<c2w_state_slot> list_gen_state_slots(const clip_hparams & hparams, const clip_model & model) {
switch (model.proj_type) {
case PROJECTOR_TYPE_QWEN3TTS_GEN: return list_c2w_state_slots(hparams, model);
case PROJECTOR_TYPE_POCKETTTS_GEN: return list_pockettts_state_slots(hparams, model);
default: return {};
}
}
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();
@ -4203,6 +4369,11 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
clip_model_loader::warmup(*ctx, *params->imgs);
}
if (params->seed != ctx->rng_seed) {
ctx->rng_seed = params->seed;
ctx->rng.seed(params->seed == UINT32_MAX ? std::random_device{}() : params->seed);
}
// build the inference graph
ggml_backend_sched_reset(ctx->sched.get());
ggml_cgraph * gf = clip_get_graph_builder(ctx, imgs, params)->build();
@ -4247,6 +4418,50 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
ggml_backend_tensor_set(cur, values.data(), 0, ggml_nbytes(cur));
};
// upload the decoder state from the previous call, or zero-fill on a cold start
auto set_gen_state_in = [&]() {
size_t offset = 0;
for (const auto & slot : list_gen_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;
}
};
// rope positions and attention mask of the mimi transformers (pocket-tts).
// the mask is causal with a sliding window, see _build_attention_mask() in the reference
auto set_pockettts_tfm_inputs = [&]() {
const int64_t n_pos = ggml_nelements(get_inp_tensor("inp_pos"));
GGML_ASSERT(n_pos > 0);
std::vector<int32_t> positions((size_t) n_pos);
for (int64_t i = 0; i < n_pos; i++) {
positions[(size_t) i] = (int32_t) i;
}
set_input_i32("inp_pos", positions);
// the preprocessor truncates the waveform to keep this mask bounded
const int64_t max_pos = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate / 120;
GGML_ASSERT(n_pos <= max_pos && "pocket-tts speaker reference too long for a dense mask");
const int64_t context = hparams.mimi_tfm_context;
std::vector<float> mask((size_t) n_pos * n_pos, -INFINITY);
for (int64_t q = 0; q < n_pos; q++) {
for (int64_t k = 0; k < n_pos; k++) {
const int64_t delta = q - k;
if (delta >= 0 && delta < context) {
mask[(size_t) q * n_pos + k] = 0.0f;
}
}
}
set_input_f32("kq_mask", mask);
};
// set input pixel values
if (!imgs.is_audio) {
size_t nelem = 0;
@ -4290,8 +4505,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
}
set_input_f32("inp_raw", inp_raw);
} 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
} else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV) {
// audio input. GEN_WAV is not here: it takes codes or feats, set in the switch below
GGML_ASSERT(imgs.entries.size() == 1);
const auto & mel_inp = imgs.entries[0];
@ -4824,6 +5039,30 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
}
set_input_i32("patches", patches);
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
{
set_pockettts_tfm_inputs();
} break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
{
if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) {
GGML_ASSERT(params->feats != nullptr);
set_input_f32("inp_feats", *params->feats);
// positions and mask are derived in-graph from the persisted counter
set_gen_state_in();
} else {
// flow matching starts from gaussian noise, std = sqrt(temp)
ggml_tensor * t = get_inp_tensor("inp_noise");
// Config.default_temperature, for a caller that does not set one
const float temp = params->temp > 0.0f ? params->temp : 0.7f;
std::normal_distribution<float> dist(0.0f, std::sqrt(temp));
std::vector<float> noise(ggml_nelements(t));
for (auto & v : noise) {
v = dist(ctx->rng);
}
set_input_f32("inp_noise", noise);
}
} break;
case PROJECTOR_TYPE_GEMMA4V:
case PROJECTOR_TYPE_GEMMA4UV:
{
@ -4948,20 +5187,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
}
}
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;
}
set_gen_state_in();
} 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];
@ -4973,11 +5199,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
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) };
std::vector<float> r = { dist(ctx->rng) };
set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r);
}
}
@ -5430,14 +5655,31 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
// for audio gen models
//
// optional outputs: a pipeline yields codes or feats, and not all have an eos head
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");
if (codes != nullptr) {
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_feats != nullptr) {
ggml_tensor * feats = ggml_graph_get_tensor(gf, "out_feats");
if (feats != nullptr) {
auto & out_feats = *params->out_feats;
out_feats.resize(ggml_nelements(feats));
ggml_backend_tensor_get(feats, out_feats.data(), 0, ggml_nbytes(feats));
}
}
if (params->out_is_eos != nullptr) {
ggml_tensor * eos = ggml_graph_get_tensor(gf, "out_eos_score");
if (eos != nullptr) {
GGML_ASSERT(ggml_nelements(eos) == 1);
float score = 0.0f;
ggml_backend_tensor_get(eos, &score, 0, sizeof(float));
*params->out_is_eos = score > hparams.gen_eos_threshold;
}
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");
@ -5449,9 +5691,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
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_codes = params->codes ? model.gen_code_head_w->ne[2] + 1 : 0;
const int64_t n_frames_w = hparams.wav_tfm_swa;
const int64_t n_frames = (int64_t) params->codes->size() / n_codes;
const int64_t n_frames = params->codes ? (int64_t) params->codes->size() / n_codes : n_frames_w;
if (n_frames < n_frames_w) {
const size_t hop = out_audio.size() / n_frames_w;
out_audio.resize((size_t) n_frames * hop);
@ -5460,12 +5702,12 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
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)) {
for (const auto & slot : list_gen_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)) {
for (const auto & slot : list_gen_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());
@ -5613,6 +5855,10 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
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_POCKETTTS_SPKENC:
return ctx->model.spk_proj_w->ne[1];
case PROJECTOR_TYPE_POCKETTTS_GEN:
return ctx->model.gen_input_lin_w->ne[1];
case PROJECTOR_TYPE_PARAKEET:
return ctx->model.mm_1_w->ne[1];
default:

5
tools/mtmd/clip.h
View file

@ -104,9 +104,14 @@ struct clip_encode_params {
int32_t top_k = 50;
float top_p = 1.0f;
std::vector<int32_t> * out_codes = nullptr; // this frame's 16 sampled codes
std::vector<float> * out_feats = nullptr; // continuous counterpart of out_codes
uint32_t seed = UINT32_MAX; // UINT32_MAX for random
float temp = 0.0f; // sampling temperature, noise scale for flow-matching decoders
bool * out_is_eos = nullptr;
// GEN_WAV
const std::vector<int32_t> * codes = nullptr; // this frame's 16 RVQ codes
const std::vector<float> * feats = nullptr; // continuous counterpart of 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

56
tools/mtmd/models/models.h
View file

@ -318,6 +318,59 @@ struct clip_graph_qwen3tts_gen : clip_graph {
};
};
//
// pocket-tts: SEANet convolution stack, shared by the voice encoder and the mimi decoder.
// stateless unless state_in is populated: convs then pad instead of carrying left-context.
//
struct clip_graph_pockettts_seanet : clip_graph {
clip_graph_pockettts_seanet(const clip_graph & parent) : clip_graph(parent) {}
ggml_cgraph * build() override { GGML_ABORT("call encode()/decode() instead"); }
// per-call streaming state, keyed by slot name (see list_pockettts_state_slots)
std::map<std::string, ggml_tensor *> state_in;
mutable std::vector<std::pair<std::string, ggml_tensor *>> state_out;
ggml_tensor * conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation,
bool pad_replicate = false, const std::string & state_name = "") const;
ggml_tensor * conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride,
const std::string & state_name = "") const;
ggml_tensor * res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation,
const std::string & state_prefix = "") const;
// x: [T, C] -> [T / hop, dim]
ggml_tensor * encode(ggml_tensor * x) const;
// x: [T, dim] -> [T * hop, 1], streams when state_in is populated
ggml_tensor * decode(ggml_tensor * x) const;
};
// mimi encoder + speaker_proj: reference waveform -> voice conditioning rows
struct clip_graph_pockettts_spkenc : clip_graph {
clip_graph_pockettts_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
ggml_cgraph * build() override;
ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const;
};
//
// pocket-tts generation:
// GEN_CODE = flow-matching decoder + end-of-speech head, one latent per call
// GEN_WAV = mimi decoder, a window of latents -> PCM
//
struct clip_graph_pockettts_gen : clip_graph {
clip_graph_pockettts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_step, int n_frames)
: clip_graph(ctx, img), gen_process(gen_process), n_step(n_step), n_frames(n_frames) {}
ggml_cgraph * build() override;
clip_gen_process_type gen_process;
int n_step; // lsd_decode steps, fixed at graph-build time
int n_frames; // GEN_WAV only: number of latents to decode
// AdaLN modulation: x * (1 + scale) + shift
ggml_tensor * modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const;
ggml_tensor * time_embed(const clip_flow_net::time_embd & te, float t) const;
ggml_tensor * flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const;
};
// one persisted state buffer used by code2wav, see qwen3tts-gen.cpp
struct c2w_state_slot {
std::string name;
@ -326,6 +379,9 @@ struct c2w_state_slot {
};
std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model);
// same, for the streaming mimi decoder (pocket-tts GEN_WAV)
std::vector<c2w_state_slot> list_pockettts_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;

291
tools/mtmd/models/pockettts-gen.cpp Normal file
View file

@ -0,0 +1,291 @@
#include "models.h"
#include <cmath>
// pocket-tts generation stages
//
// GEN_CODE: backbone hidden state -> next 32-d latent (flow matching) + end-of-speech score
// GEN_WAV : a window of latents -> PCM, through the mimi decoder
//
// there is no codebook anywhere, "codes" in the mtmd API are continuous features here
ggml_tensor * clip_graph_pockettts_gen::modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const {
ggml_tensor * cur = ggml_mul(ctx0, x, ggml_scale_bias(ctx0, scale, 1.0f, 1.0f));
return ggml_add(ctx0, cur, shift);
}
// see TimestepEmbedder in the reference
ggml_tensor * clip_graph_pockettts_gen::time_embed(const clip_flow_net::time_embd & te, float t) const {
// t is a graph-build constant, so the cos/sin table can be folded into a scaled copy
ggml_tensor * args = ggml_scale(ctx0, te.freqs, t);
ggml_tensor * emb = ggml_concat(ctx0, ggml_cos(ctx0, args), ggml_sin(ctx0, args), 0);
ggml_tensor * cur = build_mm(te.up_w, emb);
cur = ggml_add(ctx0, cur, te.up_b);
cur = ggml_silu(ctx0, cur);
cur = build_mm(te.down_w, cur);
cur = ggml_add(ctx0, cur, te.down_b);
// this "RMSNorm" divides by the unbiased variance, not the mean square
// it also rescales the input, not the centered value, see _rms_norm() in mlp.py
{
const int64_t n = cur->ne[0];
ggml_tensor * mean = ggml_mean(ctx0, cur);
ggml_tensor * dev = ggml_sub(ctx0, cur, mean);
ggml_tensor * var = ggml_mean(ctx0, ggml_sqr(ctx0, dev));
var = ggml_scale_bias(ctx0, var, (float) n / (float) (n - 1), 1e-5f);
cur = ggml_div(ctx0, cur, ggml_sqrt(ctx0, var));
cur = ggml_mul(ctx0, cur, te.norm);
}
return cur;
}
// one velocity evaluation: v(cond, s, t, x)
ggml_tensor * clip_graph_pockettts_gen::flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const {
const auto & flow = model.flow;
ggml_tensor * cur = build_mm(flow.input_proj_w, x);
cur = ggml_add(ctx0, cur, flow.input_proj_b);
// the two time conditions are averaged, then added to the projected backbone state
ggml_tensor * ts = ggml_add(ctx0, time_embed(flow.time[0], s), time_embed(flow.time[1], t));
ts = ggml_scale(ctx0, ts, 1.0f / (float) flow.time.size());
ggml_tensor * c = build_mm(flow.cond_embd_w, cond);
c = ggml_add(ctx0, c, flow.cond_embd_b);
ggml_tensor * y = ggml_add(ctx0, ts, c);
cb(y, "flow_cond", -1);
const int64_t n_ch = flow.blocks.empty() ? 0 : flow.blocks[0].norm_w->ne[0];
for (size_t il = 0; il < flow.blocks.size(); il++) {
const auto & blk = flow.blocks[il];
ggml_tensor * mod = build_mm(blk.ada_w, ggml_silu(ctx0, y));
mod = ggml_add(ctx0, mod, blk.ada_b);
ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0);
ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]);
ggml_tensor * gate = ggml_view_1d(ctx0, mod, n_ch, (size_t) 2 * n_ch * mod->nb[0]);
ggml_tensor * h = build_norm(cur, blk.norm_w, blk.norm_b, NORM_TYPE_NORMAL, 1e-6f, (int) il);
h = modulate(h, shift, scale);
h = build_mm(blk.up_w, h);
h = ggml_add(ctx0, h, blk.up_b);
h = ggml_silu(ctx0, h);
h = build_mm(blk.down_w, h);
h = ggml_add(ctx0, h, blk.down_b);
cur = ggml_add(ctx0, cur, ggml_mul(ctx0, gate, h));
cb(cur, "flow_blk", (int) il);
}
// final layer: the norm has no weights, only the AdaLN modulation
ggml_tensor * mod = build_mm(flow.final_ada_w, ggml_silu(ctx0, y));
mod = ggml_add(ctx0, mod, flow.final_ada_b);
ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0);
ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]);
cur = build_norm(cur, nullptr, nullptr, NORM_TYPE_NORMAL, 1e-6f, -1);
cur = modulate(cur, shift, scale);
cur = build_mm(flow.final_proj_w, cur);
cur = ggml_add(ctx0, cur, flow.final_proj_b);
return cur;
}
// state carried between GEN_WAV calls: rope offset, per-layer KV window, conv left context
// and the transposed-conv overlap tails
std::vector<c2w_state_slot> list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model) {
std::vector<c2w_state_slot> slots;
if (model.gen_upsample_w == nullptr) {
return slots; // not a pocket-tts decoder
}
const auto & seanet = model.seanet;
// the slots below are sized from these
GGML_ASSERT(!model.gen_tfm_layers.empty());
GGML_ASSERT((int) seanet.stages.size() >= hparams.seanet_n_stage);
GGML_ASSERT((int) hparams.seanet_ratios.size() >= hparams.seanet_n_stage);
GGML_ASSERT(hparams.mimi_tfm_context > 1 && hparams.mimi_downsample > 0);
slots.push_back({"tfm_pos", 1, 1});
const int64_t n_embd_a = model.gen_tfm_layers[0].q_w->ne[1];
const int64_t prefix = hparams.mimi_tfm_context - 1;
for (size_t il = 0; il < model.gen_tfm_layers.size(); il++) {
slots.push_back({"tfm_k_" + std::to_string(il), n_embd_a, prefix});
slots.push_back({"tfm_v_" + std::to_string(il), n_embd_a, prefix});
}
// upsample is depthwise, its output channel count is the input one
slots.push_back({"up", model.gen_upsample_w->ne[0] - hparams.mimi_downsample, model.gen_upsample_w->ne[2]});
slots.push_back({"dec_in", seanet.conv_in_w->ne[0] - 1, seanet.conv_in_w->ne[1]});
for (int i = 0; i < hparams.seanet_n_stage; i++) {
const auto & stage = seanet.stages[i];
const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i];
slots.push_back({"dec_up_" + std::to_string(i), stage.scale_conv_w->ne[0] - stride, stage.scale_conv_w->ne[1]});
slots.push_back({"dec_res_" + std::to_string(i), stage.res_conv1_w->ne[0] - 1, stage.res_conv1_w->ne[1]});
}
slots.push_back({"dec_out", seanet.conv_out_w->ne[0] - 1, seanet.conv_out_w->ne[1]});
return slots;
}
ggml_cgraph * clip_graph_pockettts_gen::build() {
if (gen_process == CLIP_GEN_PROCESS_GEN_CODE) {
// the backbone hidden state arrives as the single batch entry
ggml_tensor * h_state = build_inp_raw(1);
h_state = ggml_reshape_2d(ctx0, h_state, n_mmproj_embd, 1);
// end-of-speech probe, thresholded on the host side
ggml_tensor * eos = build_mm(model.gen_out_eos_w, h_state);
eos = ggml_add(ctx0, eos, model.gen_out_eos_b);
ggml_set_name(eos, "out_eos_score");
ggml_set_output(eos);
ggml_build_forward_expand(gf, eos);
const int64_t n_latent = model.gen_input_lin_w->ne[0];
ggml_tensor * noise = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_latent, 1);
ggml_set_name(noise, "inp_noise");
ggml_set_input(noise);
// lsd_decode: integrate the velocity field from the noise sample
ggml_tensor * cur = noise;
for (int i = 0; i < n_step; i++) {
const float s = (float) i / (float) n_step;
const float t = (float) (i + 1) / (float) n_step;
ggml_tensor * v = flow_forward(h_state, cur, s, t);
cur = ggml_add(ctx0, cur, ggml_scale(ctx0, v, 1.0f / (float) n_step));
}
cb(cur, "flow_latent", -1);
ggml_set_name(cur, "out_feats");
ggml_set_output(cur);
ggml_build_forward_expand(gf, cur);
// the same latent, projected into the backbone's input space for the next step
ggml_tensor * embd = build_mm(model.gen_input_lin_w, cur);
cb(embd, "gen_embd", -1);
ggml_build_forward_expand(gf, embd);
return gf;
}
// GEN_WAV: [32, n_frames] latents -> PCM
ggml_tensor * feats = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32,
model.gen_input_lin_w->ne[0], n_frames);
ggml_set_name(feats, "inp_feats");
ggml_set_input(feats);
// denormalize, then the DummyQuantizer up-projection
ggml_tensor * cur = ggml_add(ctx0, ggml_mul(ctx0, feats, model.gen_emb_std), model.gen_emb_mean);
cur = build_mm(model.gen_quant_out_w, cur);
cb(cur, "quant_out", -1);
clip_graph_pockettts_seanet seanet(*this);
for (const auto & slot : list_pockettts_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);
seanet.state_in[slot.name] = t;
}
// model frame rate -> encoder frame rate, depthwise transposed conv
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
cur = seanet.conv_transpose1d(cur, model.gen_upsample_w, nullptr, hparams.mimi_downsample, "up");
cb(cur, "mimi_upsample", -1);
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
// positions continue across calls, the counter lives in the state
const int64_t n_pos = cur->ne[1];
const int64_t prefix = hparams.mimi_tfm_context - 1;
const int64_t n_kv = prefix + n_pos;
ggml_tensor * base = ggml_reshape_1d(ctx0, seanet.state_in.at("tfm_pos"), 1);
ggml_tensor * inp_pos = ggml_cast(ctx0, ggml_add(ctx0, ggml_arange(ctx0, 0.0f, (float) n_pos, 1.0f), base),
GGML_TYPE_I32);
seanet.state_out.push_back({"tfm_pos", ggml_scale_bias(ctx0, seanet.state_in.at("tfm_pos"), 1.0f, (float) n_pos)});
// banded causal mask over [cached prefix | this chunk]
// the last factor masks out cache rows that hold no real frame yet
ggml_tensor * pos_k = ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) n_kv, 1.0f), n_kv, 1);
ggml_tensor * pos_q = ggml_reshape_2d(ctx0, ggml_arange(ctx0, (float) prefix, (float) (prefix + n_pos), 1.0f), 1, n_pos);
ggml_tensor * diff = ggml_sub(ctx0, ggml_repeat_4d(ctx0, pos_q, n_kv, n_pos, 1, 1), pos_k);
ggml_tensor * keep = ggml_mul(ctx0,
ggml_step(ctx0, ggml_scale_bias(ctx0, diff, 1.0f, 0.5f)), // delta >= 0
ggml_step(ctx0, ggml_scale_bias(ctx0, diff, -1.0f, (float) hparams.mimi_tfm_context - 0.5f))); // delta < context
keep = ggml_mul(ctx0, keep,
ggml_step(ctx0, ggml_scale_bias(ctx0, ggml_add(ctx0, pos_k, base), 1.0f, 0.5f - (float) prefix)));
ggml_tensor * kq_mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, keep), n_kv, n_pos, 1, 1);
for (int il = 0; il < n_layer; il++) {
const auto & layer = model.gen_tfm_layers[il];
ggml_tensor * inp = cur;
cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
ggml_tensor * Qcur = build_mm(layer.q_w, cur);
ggml_tensor * Kcur = build_mm(layer.k_w, cur);
ggml_tensor * Vcur = build_mm(layer.v_w, cur);
Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
// prepend the cached window, then keep this chunk's tail for the next call
const std::string k_name = "tfm_k_" + std::to_string(il);
const std::string v_name = "tfm_v_" + std::to_string(il);
ggml_tensor * k_full = ggml_concat(ctx0, seanet.state_in.at(k_name),
ggml_reshape_2d(ctx0, Kcur, d_head * n_head, n_pos), 1);
ggml_tensor * v_full = ggml_concat(ctx0, seanet.state_in.at(v_name), Vcur, 1);
seanet.state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, k_full, k_full->ne[0], prefix,
k_full->nb[1], (size_t) n_pos * k_full->nb[1]))});
seanet.state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, v_full, v_full->ne[0], prefix,
v_full->nb[1], (size_t) n_pos * v_full->nb[1]))});
ggml_tensor * q_cur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, n_pos, 1);
ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_full, d_head, n_head, n_kv, 1);
ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_full, d_head, n_head, n_kv, 1);
cur = build_attn(layer.o_w, nullptr, q_cur, k_cur, v_cur, kq_mask, kq_scale, il);
cur = ggml_mul(ctx0, cur, layer.ls_1_w);
cur = ggml_add(ctx0, cur, inp);
inp = cur;
cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il);
cur = ggml_mul(ctx0, cur, layer.ls_2_w);
cur = ggml_add(ctx0, cur, inp);
}
cb(cur, "mimi_dec_tfm", -1);
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
cur = seanet.decode(cur);
for (const auto & s : seanet.state_out) {
ggml_set_name(s.second, ("state_out_" + s.first).c_str());
ggml_set_output(s.second);
ggml_build_forward_expand(gf, s.second);
}
// [n_samples, 1] -> [n_samples], clamped like the reference output
cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]);
cur = ggml_clamp(ctx0, cur, -1.0f, 1.0f);
ggml_set_name(cur, "out_audio");
ggml_set_output(cur);
ggml_build_forward_expand(gf, cur);
return gf;
}

162
tools/mtmd/models/pockettts-seanet.cpp Normal file
View file

@ -0,0 +1,162 @@
#include "models.h"
// SEANet convolution stack of the mimi codec, see pocket_tts/modules/seanet.py
//
// tensors are T-first here: [T, C]
// the convs are causal: left context comes from a state slot, or from padding on a cold start
static int64_t div_ceil(int64_t a, int64_t b) {
return a / b + (a % b ? 1 : 0);
}
// x: [T, IC], w: [K, IC, OC] -> [T / stride, OC]
// the convs are causal, so the whole K - stride padding goes on the left
ggml_tensor * clip_graph_pockettts_seanet::conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation,
bool pad_replicate, const std::string & state_name) const {
const int64_t k_size = (w->ne[0] - 1) * dilation + 1;
const int64_t p_total = k_size - stride;
// trailing padding so the last frame is not dropped, see pad_for_conv1d() in conv.py
const int64_t n_frames = div_ceil(x->ne[0] - k_size + p_total, stride);
const int64_t ideal_len = n_frames * stride + k_size - p_total;
const int64_t p_extra = ideal_len - x->ne[0];
if (!state_name.empty() && p_total > 0) {
// streaming: the left context is the tail of the previous call
ggml_tensor * left = state_in.at(state_name); // [p_total, IC]
x = ggml_concat(ctx0, left, x, 0);
state_out.push_back({state_name,
ggml_cont(ctx0, ggml_view_2d(ctx0, x, p_total, x->ne[1], x->nb[1],
(size_t) (x->ne[0] - p_total) * x->nb[0]))});
} else if (pad_replicate && p_total > 0) {
// the resamplers repeat the first frame instead of zero-padding
ggml_tensor * first = ggml_view_2d(ctx0, x, 1, x->ne[1], x->nb[1], 0);
ggml_tensor * left = ggml_repeat_4d(ctx0, first, p_total, x->ne[1], 1, 1);
x = ggml_concat(ctx0, left, x, 0);
x = ggml_pad_ext(ctx0, x, 0, p_extra, 0, 0, 0, 0, 0, 0);
} else {
x = ggml_pad_ext(ctx0, x, p_total, p_extra, 0, 0, 0, 0, 0, 0);
}
ggml_tensor * y = ggml_conv_1d(ctx0, w, x, stride, 0, dilation);
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]));
}
return y;
}
// x: [T, IC], w: [K, OC/groups, IC] -> [T * stride, OC]
// the K - stride overlap tail belongs to the next call: added to its head when streaming, else dropped
ggml_tensor * clip_graph_pockettts_seanet::conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride,
const std::string & state_name) const {
const int64_t K = w->ne[0];
const int64_t T = x->ne[0];
const int64_t p_total = K - stride;
const bool depthwise = w->ne[1] == 1 && w->ne[2] > 1;
const int64_t OC = depthwise ? w->ne[2] : w->ne[1];
const int64_t emit_len = T * stride;
// one column per input step, holding the [K, OC] window that col2im scatter-adds at t * stride
ggml_tensor * col;
if (depthwise) {
// one group per channel: a batched matmul over the channels scales the kernel by each step
ggml_tensor * krn = ggml_reshape_3d(ctx0, w, 1, K, OC); // [1, K, OC]
ggml_tensor * xs = ggml_reshape_3d(ctx0, x, 1, T, OC); // [1, T, OC]
col = ggml_mul_mat(ctx0, krn, xs); // [K, T, OC]
col = ggml_cont(ctx0, ggml_permute(ctx0, col, 0, 2, 1, 3)); // [K, OC, T]
col = ggml_reshape_2d(ctx0, col, K * OC, T);
} else {
ggml_tensor * w2 = ggml_reshape_2d(ctx0, w, K * OC, w->ne[2]);
w2 = ggml_cont(ctx0, ggml_transpose(ctx0, w2)); // [IC, K * OC]
ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [IC, T]
col = ggml_mul_mat(ctx0, w2, xt);
}
ggml_tensor * full = ggml_col2im_1d(ctx0, col, stride, OC, 0); // [emit_len + p_total, OC]
ggml_tensor * out;
if (state_name.empty() || p_total == 0) {
out = ggml_cont(ctx0, ggml_view_2d(ctx0, full, emit_len, full->ne[1], full->nb[1], 0));
} else {
// overlap-add the tail the previous call held back
ggml_tensor * prev = state_in.at(state_name); // [p_total, OC]
ggml_tensor * head = ggml_add(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1], 0), prev);
if (emit_len > p_total) {
ggml_tensor * rest = ggml_view_2d(ctx0, full, emit_len - p_total, full->ne[1], full->nb[1],
(size_t) p_total * full->nb[0]);
out = ggml_concat(ctx0, head, rest, 0);
} else {
out = head;
}
state_out.push_back({state_name,
ggml_cont(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1],
(size_t) emit_len * full->nb[0]))});
}
if (b) {
out = ggml_add(ctx0, out, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
}
return out;
}
ggml_tensor * clip_graph_pockettts_seanet::res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation,
const std::string & state_prefix) const {
ggml_tensor * h = ggml_elu(ctx0, x);
h = conv1d(h, stage.res_conv1_w, stage.res_conv1_b, 1, dilation, false, state_prefix);
h = ggml_elu(ctx0, h);
// the second conv is pointwise, it needs no left context
h = conv1d(h, stage.res_conv2_w, stage.res_conv2_b, 1, 1);
return ggml_add(ctx0, x, h);
}
ggml_tensor * clip_graph_pockettts_seanet::encode(ggml_tensor * x) const {
const auto & seanet = model.seanet;
ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1);
cb(cur, "seanet_enc_in", -1);
for (int i = 0; i < hparams.seanet_n_stage; i++) {
const auto & stage = seanet.stages[i];
const int stride = hparams.seanet_ratios[i];
cur = res_unit(cur, stage, 1);
cur = ggml_elu(ctx0, cur);
cur = conv1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride, 1);
cb(cur, "seanet_enc_stage", i);
}
cur = ggml_elu(ctx0, cur);
cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1);
cb(cur, "seanet_enc_out", -1);
return cur;
}
ggml_tensor * clip_graph_pockettts_seanet::decode(ggml_tensor * x) const {
const auto & seanet = model.seanet;
const bool stream = !state_in.empty();
ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1, false,
stream ? "dec_in" : "");
cb(cur, "seanet_dec_in", -1);
for (int i = 0; i < hparams.seanet_n_stage; i++) {
const auto & stage = seanet.stages[i];
// the decoder mirrors the encoder, so the ratios are walked backwards
const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i];
const std::string id = std::to_string(i);
cur = ggml_elu(ctx0, cur);
cur = conv_transpose1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride,
stream ? "dec_up_" + id : "");
cur = res_unit(cur, stage, 1, stream ? "dec_res_" + id : "");
cb(cur, "seanet_dec_stage", i);
}
cur = ggml_elu(ctx0, cur);
cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1, false,
stream ? "dec_out" : "");
cb(cur, "seanet_dec_out", -1);
return cur;
}

77
tools/mtmd/models/pockettts-spkenc.cpp Normal file
View file

@ -0,0 +1,77 @@
#include "models.h"
// voice-prompt encoder: raw 24kHz waveform -> one conditioning row per 12.5Hz frame
// mimi encoder (SEANet + transformer + downsample), then flow_lm.speaker_proj_weight
// pre-norm block with layer scale on both residual paths, see mimi_transformer.py
ggml_tensor * clip_graph_pockettts_spkenc::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const {
ggml_tensor * inp = cur;
cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
ggml_tensor * Qcur = build_mm(layer.q_w, cur);
ggml_tensor * Kcur = build_mm(layer.k_w, cur);
ggml_tensor * Vcur = build_mm(layer.v_w, cur);
const int64_t n_pos = cur->ne[1];
Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos);
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
cur = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, kq_mask, kq_scale, il);
cur = ggml_mul(ctx0, cur, layer.ls_1_w);
cur = ggml_add(ctx0, cur, inp);
inp = cur;
cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il);
cur = ggml_mul(ctx0, cur, layer.ls_2_w);
cur = ggml_add(ctx0, cur, inp);
return cur;
}
ggml_cgraph * clip_graph_pockettts_spkenc::build() {
// the preprocessor hands over the waveform as a single-row "mel", already [n_samples, 1]
ggml_tensor * inp_raw = build_inp_raw(1);
ggml_tensor * cur = ggml_reshape_2d(ctx0, inp_raw, inp_raw->ne[0], inp_raw->ne[1]);
clip_graph_pockettts_seanet seanet(*this);
cur = seanet.encode(cur);
cb(cur, "mimi_enc", -1);
// [T, 512] -> transformer works on [512, T]
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, cur->ne[1]);
ggml_set_name(inp_pos, "inp_pos");
ggml_set_input(inp_pos);
// the mimi transformer is causal with a sliding window, see _build_attention_mask()
ggml_tensor * kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, cur->ne[1], cur->ne[1]);
ggml_set_name(kq_mask, "kq_mask");
ggml_set_input(kq_mask);
for (int il = 0; il < n_layer; il++) {
cur = tfm_layer_forward(cur, model.layers[il], inp_pos, kq_mask, il);
}
cb(cur, "mimi_enc_tfm", -1);
// downsample to the model frame rate, [512, T] -> [T, 512] -> [T / 16, 32]
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
cur = seanet.conv1d(cur, model.downsample_w, nullptr, hparams.mimi_downsample, 1, true);
cb(cur, "mimi_downsample", -1);
// voice latent -> backbone embd
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
cur = build_mm(model.spk_proj_w, cur);
cb(cur, "spk_proj", -1);
ggml_build_forward_expand(gf, cur);
return gf;
}

4
tools/mtmd/models/qwen3tts-gen.cpp
View file

@ -610,6 +610,10 @@ std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, c
const auto & c2w = model.c2w;
std::vector<c2w_state_slot> slots;
if (c2w.pre_conv_w == nullptr) {
return slots; // not a code2wav model, it keeps no state between calls
}
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)

38
tools/mtmd/mtmd-audio.cpp
View file

@ -1423,3 +1423,41 @@ std::vector<float> mtmd_audio_streaming_istft::flush() {
return output;
}
//
// mtmd_audio_preprocessor_pockettts
//
// mimi takes the raw 24kHz waveform, there is no mel front-end
// the samples are handed over as a single-row "mel", to reuse the normal chunk path
//
bool mtmd_audio_preprocessor_pockettts::preprocess(const float * samples,
size_t n_samples,
std::vector<mtmd_audio_mel> & output) {
// the encoder needs whole frames, see pad_for_conv1d() in the reference
const int64_t frame_size = (int64_t) hparams.mimi_downsample * 120;
if (n_samples == 0 || frame_size <= 0) {
return false;
}
// the mimi transformer mask is dense, so cost is quadratic in the reference length
const int64_t max_samples = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate;
if ((int64_t) n_samples > max_samples) {
LOG_WRN("%s: speaker reference is %.1f s, truncating to the first %d s\n", __func__,
(double) n_samples / hparams.audio_sample_rate, clip_hparams::pockettts_max_spk_seconds);
n_samples = (size_t) max_samples;
}
const int64_t n_frames = (int64_t) (n_samples + frame_size - 1) / frame_size;
const int64_t n_padded = n_frames * frame_size;
mtmd_audio_mel out;
out.n_mel = 1;
out.n_len = n_padded;
out.n_len_org = (int64_t) n_samples;
out.data.assign((size_t) n_padded, 0.0f);
std::copy(samples, samples + n_samples, out.data.begin());
output.push_back(std::move(out));
return true;
}

7
tools/mtmd/mtmd-audio.h
View file

@ -129,6 +129,13 @@ struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor {
mtmd_audio_cache cache;
};
// mimi convolves the waveform directly, so this only pads it to a whole number of frames
struct mtmd_audio_preprocessor_pockettts : mtmd_audio_preprocessor {
mtmd_audio_preprocessor_pockettts(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;
};
struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor {
mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { }
void initialize() override;

583
tools/mtmd/mtmd-helper-gen.cpp
View file

@ -5,6 +5,8 @@
#include "../src/llama-ext.h"
#include <algorithm>
#include <cctype>
#include <cmath>
#include <cstring>
#include <memory>
#include <string>
@ -87,7 +89,8 @@ public:
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;
// set out_stop on end-of-speech, h_state_out must be null if no frame is generated
virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) = 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:
@ -200,8 +203,10 @@ public:
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;
const mtmd_gen_inp def = mtmd_gen_inp_default(mctx);
top_k = inp->top_k > 0 ? inp->top_k : def.top_k;
top_p = inp->top_p > 0 ? inp->top_p : def.top_p;
seed = inp->seed;
out_type = inp->out_type;
// the prompt above holds the whole text stream up to tts_eos, so every generated
@ -241,13 +246,26 @@ public:
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{};
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
if (sampled == LLAMA_TOKEN_NULL) {
LOG_ERR("mtmd_helper_gen_audio: qwen3tts requires a token sampled from the backbone\n");
return 1;
}
// backbone signals end-of-speech with a token, no frame for this step
if (sampled == codec_eos || llama_vocab_is_eog(vocab, sampled)) {
*out_stop = true;
*h_state_out = nullptr;
return 0;
}
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
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;
inp.seed = seed;
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n");
@ -384,10 +402,11 @@ private:
if (codes_buf.empty()) {
return true;
}
mtmd_gen_inp inp{};
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
inp.codes = codes_buf.data();
inp.n_codes = codes_buf.size();
inp.seed = seed; // same seed as gen_code, else clip reseeds mid-generation
inp.state_data = c2w_state.empty() ? nullptr : (const char *) c2w_state.data();
inp.state_size = c2w_state.size();
mtmd_gen_out out{};
@ -427,8 +446,9 @@ private:
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;
int32_t top_k = 50;
float top_p = 1.0f;
uint32_t seed = UINT32_MAX;
std::vector<int32_t> codes_buf;
std::vector<uint8_t> c2w_state;
std::vector<float> audio_pcm;
@ -438,10 +458,547 @@ private:
std::vector<char> out_buf;
};
// settings that only live in the reference's per-pack yaml, not in the checkpoint
// the english packs share the same shapes and tokenizer, but disagree on these
// all three are 0 / false when the pack does not tune them, the model default is then used
struct pockettts_pack_settings {
float temp = 0.0f;
int frames_after_eos = 0;
bool pad_short_text = false;
};
static pockettts_pack_settings pockettts_pack(const char * variant) {
static const std::unordered_map<std::string, pockettts_pack_settings> packs = {
{ "english", { 0.3f, 0, false } },
{ "english_2026-01", { 0.7f, 0, true } },
{ "english_2026-04", { 0.3f, 0, false } },
{ "french_24l", { 0.7f, 8, false } },
};
auto it = packs.find(variant ? variant : "");
if (it == packs.end()) {
LOG_WRN("mtmd_helper_gen_audio: no tuned settings for pocket-tts variant \"%s\"\n",
variant ? variant : "");
return {};
}
return it->second;
}
// pocket-tts: the backbone emits no token, the flow net turns each hidden state into a latent
// the end-of-speech head also lives in the mmproj
class pockettts_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;
feats_buf.clear();
dec_state.clear();
audio_pcm.clear();
h_state_buf.clear();
out_buf.clear();
prompt_embd_buf.clear();
prompt_batch.reset();
n_prompt = 0;
prompt_pos = 0;
step_idx = 0;
eos_step = -1;
chunks.clear();
chunk_idx = 0;
n_voice_pos = 0;
chunk_budget = 0;
}
int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
reset();
seq_id = inp->seq_id;
if (!ensure_cache()) {
return 1;
}
std::vector<float> voice;
if (inp->speaker_ref) {
if (!encode_speaker(inp->speaker_ref, voice)) {
return 1;
}
}
pack = pockettts_pack(info.model_variant);
const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len),
pack.pad_short_text);
if (text.empty()) {
LOG_ERR("mtmd_helper_gen_audio: empty prompt\n");
return 1;
}
std::vector<llama_token> ids(text.size() + 16);
int n_ids = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(), ids.data(),
(int32_t) ids.size(), false, false);
if (n_ids <= 0) {
LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
return 1;
}
ids.resize((size_t) n_ids);
// long inputs degrade badly, so each chunk restarts from the voice conditioning
// see split_into_best_sentences() in the reference
chunks = split_chunks(ids);
chunk_idx = 0;
if (chunks.size() > 1) {
LOG_INF("mtmd_helper_gen_audio: %d tokens split into %zu chunks\n", n_ids, chunks.size());
}
const int n_e = n_embd;
// sequence order is voice, then text, then the audio BOS that starts generation
if (!voice.empty()) {
GGML_ASSERT(voice.size() % (size_t) n_e == 0);
if (bos_before_voice != LLAMA_TOKEN_NULL) {
push_embd_row(prompt_embd_buf, bos_before_voice);
}
prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end());
}
// every later chunk rewinds to here and re-prompts, so the voice stays primed
n_voice_pos = (int) (prompt_embd_buf.size() / (size_t) n_e);
for (llama_token t : chunks[0]) {
push_embd_row(prompt_embd_buf, t);
}
push_embd_row(prompt_embd_buf, audio_bos);
arm_chunk_budget(0);
n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e);
prompt_batch.reset(new decode_embd_batch(prompt_embd_buf.data(), n_prompt, 1, n_e));
prompt_batch->set_position_normal(0, seq_id);
prompt_pos = 0;
seed = inp->seed;
out_type = inp->out_type;
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);
if ((prompt_pos + n_tokens_batch) == n_prompt) {
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_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, bool * out_stop) override {
(void) sampled; // the backbone output is continuous, there is no token to consume
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
inp.embd = const_cast<float *>(h_state_in);
// clip only reseeds when the seed changes, so pass the same one on every step
inp.seed = seed;
if (pack.temp > 0.0f) {
inp.temp = pack.temp;
}
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: flow decode failed\n");
return 1;
}
if (out.is_eos && eos_step < 0) {
eos_step = step_idx;
}
// the frame of the stopping step is discarded, matching _autoregressive_generation().
// the budget is the reference's fallback for a chunk whose eos head never fires
const bool chunk_done = (eos_step >= 0 && step_idx >= eos_step + frames_after_eos) ||
step_idx >= chunk_budget;
if (chunk_done) {
if (eos_step < 0) {
LOG_WRN("mtmd_helper_gen_audio: chunk %zu hit its budget without end-of-speech\n", chunk_idx);
}
return finish_chunk(h_state_out, out_stop);
}
feats_buf.insert(feats_buf.end(), out.feats, out.feats + out.n_feats);
step_idx++;
if (out.n_feats > 0 && feats_buf.size() / out.n_feats >= window_frames) {
if (!flush_gen_wav()) {
return 1;
}
}
decode_embd_batch batch_embd(const_cast<float *>(out.embd), 1, 1, n_embd);
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;
}
// bos_before_voice is optional, some packs do not insert it
bos_before_voice = find_special_token(vocab, "<|bos_before_voice|>");
audio_bos = find_special_token(vocab, "<|audio_bos|>");
if (audio_bos == LLAMA_TOKEN_NULL) {
LOG_ERR("mtmd_helper_gen_audio: missing <|audio_bos|> 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;
}
GGML_ASSERT(n_embd > 0 && n_tok_embd % (uint32_t) n_embd == 0);
specials_ok = true;
return true;
}
// the table can be shorter than the vocab, so bound the row lookup
void push_embd_row(std::vector<float> & dst, llama_token t) const {
const size_t n_rows = tok_embd.size() / (size_t) n_embd;
GGML_ASSERT(t >= 0 && (size_t) t < n_rows);
dst.insert(dst.end(),
tok_embd.begin() + (size_t) t * n_embd,
tok_embd.begin() + (size_t) (t + 1) * n_embd);
}
// token ids of the pieces the reference splits on, see split_into_best_sentences().
// the leading token is dropped, it is the tokenizer's dummy prefix
std::vector<llama_token> punct_ids(const char * s) const {
std::vector<llama_token> ids(16);
const int n = llama_tokenize(vocab, s, (int32_t) strlen(s), ids.data(), (int32_t) ids.size(), false, false);
if (n <= 1) {
return {};
}
return std::vector<llama_token>(ids.begin() + 1, ids.begin() + n);
}
// cut after runs of boundary tokens, so punctuation stays with the sentence it ends
static std::vector<std::vector<llama_token>> split_on(const std::vector<llama_token> & ids,
const std::vector<llama_token> & boundary) {
std::vector<std::vector<llama_token>> out;
size_t start = 0;
bool prev_was_boundary = false;
for (size_t i = 0; i < ids.size(); i++) {
const bool is_boundary = std::find(boundary.begin(), boundary.end(), ids[i]) != boundary.end();
if (!is_boundary && prev_was_boundary) {
out.emplace_back(ids.begin() + start, ids.begin() + i);
start = i;
}
prev_was_boundary = is_boundary;
}
out.emplace_back(ids.begin() + start, ids.end());
return out;
}
std::vector<std::vector<llama_token>> split_chunks(const std::vector<llama_token> & ids) const {
if ((int) ids.size() <= max_chunk_tokens) {
return { ids };
}
const std::vector<llama_token> eos_punct = punct_ids(".!...?");
const std::vector<llama_token> mid_punct = punct_ids(",;:");
// oversized sentences are split again on weaker punctuation, else words get skipped
std::vector<std::vector<llama_token>> segments;
for (auto & seg : split_on(ids, eos_punct)) {
if ((int) seg.size() <= max_chunk_tokens) {
segments.push_back(std::move(seg));
continue;
}
auto sub = split_on(seg, mid_punct);
if (sub.size() > 1) {
for (auto & s : sub) {
segments.push_back(std::move(s));
}
} else {
segments.push_back(std::move(seg));
}
}
std::vector<std::vector<llama_token>> out;
for (auto & seg : segments) {
if (seg.empty()) {
continue;
}
if (!out.empty() && (int) (out.back().size() + seg.size()) <= max_chunk_tokens) {
out.back().insert(out.back().end(), seg.begin(), seg.end());
} else {
out.push_back(std::move(seg));
}
}
if (out.empty()) {
out.push_back(ids);
}
for (const auto & c : out) {
if ((int) c.size() > max_chunk_tokens) {
LOG_WRN("mtmd_helper_gen_audio: chunk of %zu tokens exceeds the %d token budget, "
"generation may skip words\n", c.size(), max_chunk_tokens);
}
}
return out;
}
// _estimate_max_gen_len() plus the per-chunk tail guess, both in frames
void arm_chunk_budget(size_t idx) {
const int n_tok = (int) chunks[idx].size();
chunk_budget = (int) std::ceil((n_tok / 3.0 + 2.0) * frame_rate);
// the pack may pin the tail, else the reference guesses it from the word count
frames_after_eos = pack.frames_after_eos > 0 ? pack.frames_after_eos : (n_tok <= 6 ? 5 : 3);
step_idx = 0;
eos_step = -1;
}
// ends the current chunk and, if there is another, re-prompts it on top of the voice
int32_t finish_chunk(const float ** h_state_out, bool * out_stop) {
if (!flush_gen_wav()) {
return 1;
}
// the decoder restarts too, the next chunk's audio is not continuous with this one
dec_state.clear();
if (chunk_idx + 1 >= chunks.size()) {
*out_stop = true;
*h_state_out = nullptr;
return 0;
}
chunk_idx++;
// drop this chunk's text and audio, keep the voice conditioning
llama_memory_seq_rm(llama_get_memory(lctx), seq_id, n_voice_pos, -1);
pos = n_voice_pos;
const int n_e = n_embd;
prompt_embd_buf.clear();
for (llama_token t : chunks[chunk_idx]) {
push_embd_row(prompt_embd_buf, t);
}
push_embd_row(prompt_embd_buf, audio_bos);
arm_chunk_budget(chunk_idx);
const int n_rows = (int) (prompt_embd_buf.size() / (size_t) n_e);
GGML_ASSERT(n_rows > 0);
decode_embd_batch batch(prompt_embd_buf.data(), n_rows, 1, n_e);
batch.set_position_normal(pos, seq_id);
batch.batch.logits[n_rows - 1] = 1;
if (llama_decode(lctx, batch.batch) != 0) {
LOG_ERR("mtmd_helper_gen_audio: chunk prompt decode failed\n");
return 1;
}
pos += n_rows;
prompt_embd_buf.clear();
const float * he = llama_get_embeddings_ith(lctx, -1);
h_state_buf.assign(he, he + n_embd);
*h_state_out = h_state_buf.data();
*out_stop = false;
return 0;
}
// same normalization as prepare_text_prompt() in the reference, it affects quality
static std::string prepare_text(const std::string & in, bool pad_short) {
std::string s;
s.reserve(in.size() + 1);
for (char c : in) {
if (c == '\n' || c == '\r') {
s += ' ';
} else if (c == ';') {
s += ',';
} else {
s += c;
}
}
const size_t b = s.find_first_not_of(' ');
const size_t e = s.find_last_not_of(' ');
if (b == std::string::npos) {
return "";
}
s = s.substr(b, e - b + 1);
if (s[0] >= 'a' && s[0] <= 'z') {
s[0] = (char) (s[0] - 'a' + 'A');
}
const unsigned char last = (unsigned char) s.back();
if (std::isalnum(last)) {
s += '.';
}
if (pad_short && count_words(s) < 5) {
s = std::string(8, ' ') + s;
}
return s;
}
static int count_words(const std::string & s) {
int n = 0;
bool in_word = false;
for (char c : s) {
if (c == ' ') {
in_word = false;
} else if (!in_word) {
in_word = true;
n++;
}
}
return n;
}
// runs the reference wav through the mimi encoder, returns one row per 12.5Hz frame
bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
if (!mtmd_support_audio(mctx)) {
LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice 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: voice 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;
}
// decodes the buffered latents, the mimi decoder state carries over between calls
bool flush_gen_wav() {
if (feats_buf.empty()) {
return true;
}
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
inp.feats = feats_buf.data();
inp.n_feats = feats_buf.size();
inp.seed = seed;
inp.state_data = dec_state.empty() ? nullptr : (const char *) dec_state.data();
inp.state_size = dec_state.size();
mtmd_gen_out out{};
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
LOG_ERR("mtmd_helper_gen_audio: mimi decode failed\n");
return false;
}
audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
dec_state.assign(out.state_data, out.state_data + out.state_size);
feats_buf.clear();
return true;
}
pockettts_pack_settings pack;
bool specials_ok = false;
llama_token bos_before_voice = LLAMA_TOKEN_NULL;
llama_token audio_bos = LLAMA_TOKEN_NULL;
std::vector<float> tok_embd;
llama_seq_id seq_id = 0;
int pos = 0;
std::vector<float> prompt_embd_buf;
std::unique_ptr<decode_embd_batch> prompt_batch;
int n_prompt = 0;
int prompt_pos = 0;
uint32_t seed = UINT32_MAX;
// end-of-speech is latched, then a few more frames are generated as tail padding
int step_idx = 0;
int eos_step = -1;
int frames_after_eos = 3;
static constexpr int max_chunk_tokens = 50; // MAX_TOKEN_PER_CHUNK in the reference
static constexpr double frame_rate = 12.5;
std::vector<std::vector<llama_token>> chunks;
size_t chunk_idx = 0;
int n_voice_pos = 0; // KV positions held by the voice conditioning
int chunk_budget = 0;
// latents are decoded a window at a time, the decoder state bridges the windows
size_t window_frames = 8;
std::vector<float> feats_buf;
std::vector<uint8_t> dec_state;
std::vector<float> audio_pcm;
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));
case MTMD_GEN_AUDIO_TYPE_POCKETTTS:
return std::unique_ptr<mtmd_gen_audio_pipeline>(new pockettts_gen_audio_pipeline(lctx, mctx));
default:
return nullptr;
}
@ -483,11 +1040,17 @@ int32_t mtmd_helper_gen_audio_step_prompt(mtmd_helper_gen_audio * ctx, int32_t n
}
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) {
const float * h_state_in, const float ** h_state_out,
bool * out_stop) {
if (!ctx->pipeline) {
return 1;
}
return ctx->pipeline->step_gen(sampled, h_state_in, h_state_out);
bool stop = false;
const int32_t ret = ctx->pipeline->step_gen(sampled, h_state_in, h_state_out, &stop);
if (out_stop) {
*out_stop = stop;
}
return ret;
}
int32_t mtmd_helper_gen_audio_get_output(mtmd_helper_gen_audio * ctx, int32_t * out_sample_rate,

16
tools/mtmd/mtmd-helper.h
View file

@ -183,8 +183,9 @@ struct mtmd_helper_gen_audio_inp {
mtmd_bitmap * speaker_ref; // optional, can be NULL
const char * lang; // optional, can be NULL
int32_t top_k;
float top_p;
int32_t top_k;
float top_p;
uint32_t seed; // UINT32_MAX for random (default: random)
enum mtmd_helper_gen_audio_outtype out_type;
};
@ -208,12 +209,15 @@ MTMD_API int32_t mtmd_helper_gen_audio_step_prompt(
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
// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token
// out_stop (optional) is set on end-of-speech, the caller must then stop the loop
// h_state_out is valid until next step_gen() or reset() call, null if no frame is generated
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);
const float ** h_state_out,
bool * out_stop);
// out_data valid until next get_output() or reset() call
// out_n_samples (optional, can be NULL) receives the number of generated PCM samples
@ -261,8 +265,8 @@ struct gen_audio {
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 step_gen(llama_token sampled, const float * h_state, const float ** h_state_out, bool * out_stop = nullptr) {
return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out, out_stop);
}
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);

93
tools/mtmd/mtmd.cpp
View file

@ -477,6 +477,7 @@ struct mtmd_context {
// 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_feats; // this frame's continuous features, if any (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
@ -979,6 +980,10 @@ struct mtmd_context {
{
audio_preproc = std::make_unique<mtmd_audio_preprocessor_qwen3tts_spk>(ctx_a);
} break;
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
{
audio_preproc = std::make_unique<mtmd_audio_preprocessor_pockettts>(ctx_a);
} break;
default:
throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj));
}
@ -1798,16 +1803,22 @@ float * mtmd_get_output_embd(mtmd_context * ctx) {
//
mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
mtmd_gen_audio_info info;
mtmd_gen_audio_info info{};
info.model_variant = "";
if (!ctx->ctx_gen_a) {
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
return info;
}
info.model_variant = clip_get_hparams(ctx->ctx_gen_a)->gen_model_variant.c_str();
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;
case PROJECTOR_TYPE_POCKETTTS_GEN:
info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS;
info.sample_rate = 24000;
break;
default:
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
break;
@ -1815,6 +1826,33 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
return info;
}
mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) {
mtmd_gen_inp inp{};
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
inp.seed = UINT32_MAX;
if (!ctx->ctx_gen_a) {
return inp;
}
switch (clip_get_projector_type(ctx->ctx_gen_a)) {
case PROJECTOR_TYPE_QWEN3TTS_GEN:
// https://huggingface.co/Qwen/Qwen3-TTS-12Hz-1.7B-Base/blob/main/generation_config.json
inp.top_k = 50;
inp.top_p = 1.0f;
inp.temp = 0.9f; // TODO: handle this on graph
break;
case PROJECTOR_TYPE_POCKETTTS_GEN:
// https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/default_parameters.py
inp.top_k = 50;
inp.top_p = 1.0f;
inp.temp = 0.7f;
break;
default:
break;
}
return inp;
}
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) {
@ -1822,6 +1860,8 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
return 1;
}
*out = {};
if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) {
const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip);
@ -1835,16 +1875,22 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
std::vector<float> out_embd(n_embd);
std::vector<int32_t> out_codes;
std::vector<float> out_feats;
bool is_eos = false;
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;
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.out_feats = &out_feats;
params.code0 = inp->code0;
params.top_k = inp->top_k;
params.top_p = inp->top_p;
params.seed = inp->seed;
params.temp = inp->temp;
params.out_is_eos = &is_eos;
if (!clip_encode(ctx_clip, &params)) {
LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__);
@ -1853,19 +1899,31 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
ctx->gen_out_embd = std::move(out_embd);
ctx->gen_out_codes = std::move(out_codes);
ctx->gen_out_feats = std::move(out_feats);
out->embd = ctx->gen_out_embd.data();
out->codes = ctx->gen_out_codes.data();
out->n_codes = ctx->gen_out_codes.size();
out->embd = ctx->gen_out_embd.data();
out->codes = ctx->gen_out_codes.data();
out->n_codes = ctx->gen_out_codes.size();
out->feats = ctx->gen_out_feats.data();
out->n_feats = ctx->gen_out_feats.size();
out->is_eos = is_eos;
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__);
const bool has_codes = inp->codes && inp->n_codes > 0;
const bool has_feats = inp->feats && inp->n_feats > 0;
if (has_codes == has_feats) {
LOG_ERR("%s: gen_wav requires exactly one of codes or feats\n", __func__);
return 1;
}
std::vector<int32_t> in_codes(inp->codes, inp->codes + inp->n_codes);
std::vector<int32_t> in_codes;
std::vector<float> in_feats;
if (has_codes) {
in_codes.assign(inp->codes, inp->codes + inp->n_codes);
} else {
in_feats.assign(inp->feats, inp->feats + inp->n_feats);
}
std::vector<uint8_t> in_state;
if (inp->state_data) {
in_state.assign(inp->state_data, inp->state_data + inp->state_size);
@ -1885,7 +1943,10 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
params.imgs = &batch;
params.n_threads = ctx->n_threads;
params.gen_process = CLIP_GEN_PROCESS_GEN_WAV;
params.codes = &in_codes;
// gen_wav draws no randomness, but keep the seed so it does not reseed mid-generation
params.seed = inp->seed;
params.codes = has_codes ? &in_codes : nullptr;
params.feats = has_feats ? &in_feats : nullptr;
params.out_audio = &ctx->gen_out_audio;
params.state_in = inp->state_data ? &in_state : nullptr;
params.state_out = &ctx->gen_out_state;

22
tools/mtmd/mtmd.h
View file

@ -344,18 +344,25 @@ MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname);
enum mtmd_gen_audio_type {
MTMD_GEN_AUDIO_TYPE_NONE, // not supported
MTMD_GEN_AUDIO_TYPE_QWEN3TTS,
MTMD_GEN_AUDIO_TYPE_POCKETTTS,
};
struct mtmd_gen_audio_info {
enum mtmd_gen_audio_type type;
int32_t sample_rate; // in Hz, for example 24000 for qwen3tts
const char * model_variant; // name of the weight variant, can be nullptr if not applicable
};
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
// for pocket-tts, this is mimi decoder
};
struct mtmd_gen_inp {
enum mtmd_gen_process_type type;
@ -364,21 +371,30 @@ struct mtmd_gen_inp {
float * embd; // the hidden state from backbone, must have n_text_embd elements
int32_t top_k;
float top_p;
uint32_t seed; // UINT32_MAX for random
float temp; // sampling temperature, or noise scale for flow-matching decoders
// for MTMD_GEN_PROCESS_TYPE_GEN_WAV
// pass either codes (discrete) or feats (continuous), depending on the pipeline
int32_t * codes;
size_t n_codes;
const float * feats;
size_t n_feats;
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;
size_t n_codes;
const float * feats; // continuous counterpart of codes
size_t n_feats;
const float * embd; // the generated hidden state, to be fed back to backbone
// it must have n_text_embd elements
bool is_eos; // only set by pipelines having the EOS head inside mmproj
// for MTMD_GEN_PROCESS_TYPE_GEN_WAV
const float * audio;
@ -386,6 +402,10 @@ struct mtmd_gen_out {
const char * state_data;
size_t state_size;
};
// defaults tuned for the loaded pipeline, callers override only what they care about
MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx);
// 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,

8
tools/mtmd/requirements.txt
View file

@ -2,11 +2,5 @@
--extra-index-url https://download.pytorch.org/whl/cpu
pillow~=11.3.0
## Embedding Gemma requires PyTorch 2.6.0 or later, bumped to 2.11.0 for compatibility
torch==2.11.0; platform_machine != "s390x" # check_requirements: ignore "=="
torch==2.11.0 # check_requirements: ignore "=="
torchvision==0.26.0; platform_machine != "s390x" # check_requirements: ignore "=="
# torch s390x packages can only be found from nightly builds
--extra-index-url https://download.pytorch.org/whl/nightly
torch>=0.0.0.dev0; platform_machine == "s390x" # check_requirements: ignore "=="
torchvision>=0.0.0.dev0; platform_machine == "s390x" # check_requirements: ignore "=="

7
tools/server/server-context.cpp
View file

@ -397,12 +397,7 @@ struct server_slot {
bool need_embd() const {
GGML_ASSERT(task);
return task->need_embd() || (spec && common_speculative_need_embd(spec));
}
bool need_embd_nextn() const {
GGML_ASSERT(task);
return spec && common_speculative_need_embd_nextn(spec);
return task->need_embd();
}
// if the context does not have a memory module then all embeddings have to be computed within a single ubatch

30
tools/server/server-tools.cpp
View file

@ -1,6 +1,7 @@
#include "server-tools.h"
#include "subproc.h"
#include "base64.hpp"
#include <filesystem>
#include <fstream>
@ -864,6 +865,7 @@ static bool path_glob_match(const std::string & pattern, const std::string & rel
//
static constexpr size_t SERVER_TOOL_READ_FILE_MAX_SIZE = 16 * 1024; // 16 KB
static constexpr size_t SERVER_TOOL_READ_FILE_MAX_SIZE_BASE64 = 32 * 1024 * 1024; // 32 MB
struct server_tool_read_file : server_tool {
server_tool_read_file() {
@ -899,6 +901,8 @@ struct server_tool_read_file : server_tool {
int start_line = json_value(params, "start_line", 1);
int end_line = json_value(params, "end_line", -1); // -1 = no limit
bool append_loc = json_value(params, "append_loc", false);
// comes from the x-resp-type header, the model cannot ask for it
bool as_base64 = json_value(params, "resp_type", std::string()) == "base64";
auto io = make_tools_io(params);
@ -906,6 +910,23 @@ struct server_tool_read_file : server_tool {
if (!io->file_size(path, file_size)) {
return {{"error", "cannot stat file: " + path}};
}
if (as_base64) {
if (file_size > SERVER_TOOL_READ_FILE_MAX_SIZE_BASE64) {
return {{"error", string_format(
"file too large (%zu bytes, max %zu)",
(size_t)file_size, SERVER_TOOL_READ_FILE_MAX_SIZE_BASE64)}};
}
std::string content;
if (!io->read_file(path, content)) {
return {{"error", "failed to open file: " + path}};
}
return {
{"base64", base64::encode(content.data(), content.size())},
{"size_bytes", (size_t) content.size()},
};
}
if (file_size > SERVER_TOOL_READ_FILE_MAX_SIZE && end_line == -1) {
return {{"error", string_format(
"file too large (%zu bytes, max %zu). Use start_line/end_line to read a portion.",
@ -2135,6 +2156,15 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools,
params["runtime"] = runtime->spec();
}
// x-resp-type header is only used by read_file for now
if (params.contains("resp_type")) {
params.erase("resp_type");
}
auto resp_type = get_header(req.headers, "x-resp-type");
if (!resp_type.empty()) {
params["resp_type"] = resp_type;
}
server_tool & tool = find_tool(tools, tool_name, stream);
if (stream) {

6
tools/server/tests/unit/test_speculative.py
View file

@ -27,8 +27,8 @@ def test_with_and_without_draft():
global server
request = {
"prompt": "I believe the meaning of life is",
"temperature": 0.8,
"top_k": 40,
"temperature": 0.2,
"top_k": 5,
"seed": 4242,
"n_predict": 16,
"return_tokens": True,
@ -36,7 +36,6 @@ def test_with_and_without_draft():
server.model_draft = None # disable draft model
server.spec_type = None
server.backend_sampling = True
server.start()
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200
@ -45,7 +44,6 @@ def test_with_and_without_draft():
# create new server with draft model
create_server()
server.backend_sampling = True
server.start()
res = server.make_request("POST", "/completion", data=request)
assert res.status_code == 200

25
tools/tts/README.md
View file

@ -32,3 +32,28 @@ llama-tts -hf ggml-org/Qwen3-TTS-12Hz-1.7B-Base-GGUF \
--tts-speaker-file speaker.mp3 \
--output out.wav
```
## Pocket TTS
Available params:
- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3). It is required, the model produces almost no audio without it
- Note: `lang` is not used, the language is a property of the weights
Example usage:
```sh
llama-tts -m pocket-tts.gguf \
-mm mmproj-pocket-tts.gguf \
-p "Hello world" \
--tts-speaker-file speaker.mp3 \
--output out.wav
```
**Note for GGUF conversion:**
The [upstream repository](https://huggingface.co/kyutai/pocket-tts) holds one complete model per language under `languages/`, next to a set of shared files at the root. Convert one of the `languages/<name>` directories, **not** the root directory:
```sh
python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --outfile pocket-tts.gguf
python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --outfile mmproj-pocket-tts.gguf
```

15
tools/tts/tts.cpp
View file

@ -119,6 +119,7 @@ int main(int argc, char ** argv) {
inp.lang = params.tts_lang.c_str();
inp.top_k = params.sampling.top_k;
inp.top_p = params.sampling.top_p;
inp.seed = params.sampling.seed;
inp.out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
//
@ -143,8 +144,7 @@ int main(int argc, char ** argv) {
}
}
const llama_vocab * vocab = llama_model_get_vocab(model);
// note: some pipelines ignore this token and use the hidden state instead
auto sample_semantic_code = [&]() -> llama_token {
llama_token t = common_sampler_sample(smpl, lctx, -1);
common_sampler_accept(smpl, t, true);
@ -159,19 +159,24 @@ int main(int argc, char ** argv) {
tts_timings timings;
const int64_t t_gen_start_us = ggml_time_us();
for (; n_frames < max_new && !llama_vocab_is_eog(vocab, sampled); n_frames++) {
bool stop = false;
while (!stop && n_frames < max_new) {
const float * h_next = nullptr;
// stage 2+3: semantic --> acoustic details --> audio waveform
// step_gen() runs both stages and returns new h_state for next step
if (gen.step_gen(sampled, h_state, &h_next) != 0) {
if (gen.step_gen(sampled, h_state, &h_next, &stop) != 0) {
LOG_ERR("step_gen failed at frame %d\n", n_frames);
return 1;
}
if (!h_next) {
break; // stopped without generating a frame
}
n_frames++;
h_state = h_next;
sampled = sample_semantic_code();
timings.report(n_frames + 1);
timings.report(n_frames);
}
const double t_gen_s = (ggml_time_us() - t_gen_start_us) / 1e6;

9
tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte
View file

@ -3,6 +3,7 @@
import { ChevronDown } from '@lucide/svelte';
import * as Collapsible from '$lib/components/ui/collapsible';
import { STATS_UNITS } from '$lib/constants';
import { gaugePopup } from '$lib/stores/context-gauge-popup.svelte';
interface Props {
currentRead: number;
@ -30,19 +31,19 @@
transientDetails
}: Props = $props();
let open = $state(false);
const hasCumulative = $derived(cumulativeRead > 0 || cumulativeOutput > 0);
const hasCurrent = $derived(currentRead > 0 || currentOutput > 0);
</script>
<Collapsible.Root bind:open class="mt-3 border-t border-border/50 pt-4">
<Collapsible.Root bind:open={gaugePopup.detailsOpen} class="mt-3 border-t border-border/50 pt-4">
<Collapsible.Trigger
class="flex w-full cursor-pointer items-center gap-1 text-xs text-muted-foreground hover:text-foreground"
>
<span>Token usage details</span>
<ChevronDown class={'ml-auto h-3 w-3 transition-transform' + (open ? ' rotate-180' : '')} />
<ChevronDown
class={'ml-auto h-3 w-3 transition-transform' + (gaugePopup.detailsOpen ? ' rotate-180' : '')}
/>
</Collapsible.Trigger>
<Collapsible.Content class="flex flex-col gap-4 text-xs pt-4">

3
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte
View file

@ -7,6 +7,7 @@
import ChatMessageToolCallBlockGetInfo from './ChatMessageToolCallBlockGetInfo.svelte';
import ChatMessageToolCallBlockGrepSearch from './ChatMessageToolCallBlockGrepSearch.svelte';
import ChatMessageToolCallBlockReadFile from './ChatMessageToolCallBlockReadFile.svelte';
import ChatMessageToolCallBlockReadMedia from './ChatMessageToolCallBlockReadMedia.svelte';
import ChatMessageToolCallBlockRunJavascript from './ChatMessageToolCallBlockRunJavascript.svelte';
import ChatMessageToolCallBlockSearchResults from './ChatMessageToolCallBlockSearchResults.svelte';
import ChatMessageToolCallBlockWriteFile from './ChatMessageToolCallBlockWriteFile.svelte';
@ -45,6 +46,8 @@
<ChatMessageToolCallBlockGetInfo {section} {isStreaming} />
{:else if section.toolName === BuiltInTool.READ_FILE}
<ChatMessageToolCallBlockReadFile {section} {open} {isStreaming} {onToggle} />
{:else if section.toolName === BuiltInTool.READ_MEDIA}
<ChatMessageToolCallBlockReadMedia {section} {open} {isStreaming} {onToggle} />
{:else if section.toolName === BuiltInTool.EDIT_FILE}
<ChatMessageToolCallBlockEditFile {section} {open} {isStreaming} {onToggle} />
{:else if section.toolName === BuiltInTool.WRITE_FILE}

37
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte
View file

@ -8,14 +8,16 @@
import { MarkdownContent, SyntaxHighlightedCode } from '$lib/components/app';
import { MAX_HEIGHT_CODE_BLOCK } from '$lib/constants';
import { getBuiltinToolUi } from '$lib/constants/built-in-tools';
import { FileTypeText, ToolResultKind } from '$lib/enums';
import { AttachmentType, FileTypeText, MimeTypeAudio, ToolResultKind } from '$lib/enums';
import type { DatabaseMessageExtra } from '$lib/types';
import {
type AgenticSection,
classifyToolResult,
formatJsonPretty,
parseToolResultWithImages
parseToolResultWithMedia,
type ToolResultLine
} from '$lib/utils';
import { createBase64DataUrl } from '$lib/utils/data-url';
interface Props {
section: AgenticSection;
@ -29,8 +31,8 @@
const title = $derived(getBuiltinToolUi(section.toolName)?.label ?? section.toolName ?? '');
const outputKind = $derived(classifyToolResult(section.toolResult));
const parsedLines = $derived(
section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : []
const parsedLines: ToolResultLine[] = $derived(
section.toolResult ? parseToolResultWithMedia(section.toolResult, attachments) : []
);
</script>
@ -103,13 +105,26 @@
<div class="font-mono text-[11px] leading-relaxed whitespace-pre-wrap">
{line.text}
</div>
{#if line.image}
<img
src={line.image.base64Url}
alt={line.image.name}
class="mt-2 mb-2 h-auto max-w-full rounded-lg"
loading="lazy"
/>
{#if line.media}
{#if line.media.type === AttachmentType.AUDIO}
{@const audioMimeType = line.media.mimeType ?? MimeTypeAudio.MP3_MPEG}
<div class="mt-2 mb-2">
<audio controls class="w-full rounded-lg">
<source
src={createBase64DataUrl(audioMimeType, line.media.base64Data)}
type={audioMimeType}
/>
Your browser does not support the audio element.
</audio>
</div>
{:else}
<img
src={line.media.base64Url}
alt={line.media.name}
class="mt-2 mb-2 h-auto max-w-full rounded-lg"
loading="lazy"
/>
{/if}
{/if}
{/each}
</div>

10
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte
View file

@ -23,7 +23,7 @@
isExitCodeSummaryLine,
parseExecShellCommandError,
parseExecShellCommandExitStatus,
parseToolResultWithImages,
parseToolResultWithMedia,
type ToolResultLine
} from '$lib/utils';
@ -53,7 +53,7 @@
);
const parsedLines: ToolResultLine[] = $derived(
section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : []
section.toolResult ? parseToolResultWithMedia(section.toolResult, attachments) : []
);
// Drop the trailing "[exit code: N]" line - rendered as a colored
@ -223,10 +223,10 @@
>
{#each outputLines as line, i (i)}
<div class="font-mono text-[11px] leading-relaxed whitespace-pre-wrap">{line.text}</div>
{#if line.image}
{#if line.media}
<img
src={line.image.base64Url}
alt={line.image.name}
src={line.media.base64Url}
alt={line.media.name}
class="mt-2 mb-2 h-auto max-w-full rounded-lg"
loading="lazy"
/>

99
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockReadMedia.svelte Normal file
View file

@ -0,0 +1,99 @@
<script lang="ts">
import { parseReadMediaMeta } from './parsers/read-media';
import ToolCallBlock from './ToolCallBlock.svelte';
import { ATTACHMENT_SAVED_REGEX } from '$lib/constants/agentic';
import { AttachmentType, MimeTypeAudio } from '$lib/enums';
import type { DatabaseMessageExtraAudioFile, DatabaseMessageExtraImageFile } from '$lib/types';
import { type AgenticSection } from '$lib/utils';
import { createBase64DataUrl } from '$lib/utils/data-url';
interface Props {
section: AgenticSection;
open: boolean;
isStreaming: boolean;
onToggle?: () => void;
}
let { isStreaming, onToggle, open, section }: Props = $props();
const readMediaMeta = $derived(parseReadMediaMeta(section));
// extractBase64Attachments swapped the data URI line for [Attachment saved: name]
// and moved the bytes to the message extras, so the name is the only link back
const mediaAttachment = $derived.by(() => {
const extras = section.toolResultExtras;
if (!extras || extras.length === 0) return null;
const match = section.toolResult?.match(ATTACHMENT_SAVED_REGEX);
if (!match) return null;
const attachmentName = match[1];
return (
extras.find(
(e): e is DatabaseMessageExtraImageFile | DatabaseMessageExtraAudioFile =>
(e.type === AttachmentType.IMAGE || e.type === AttachmentType.AUDIO) &&
e.name === attachmentName
) ?? null
);
});
const audioMimeType = $derived(readMediaMeta?.mimeType ?? MimeTypeAudio.MP3_MPEG);
</script>
<ToolCallBlock {section} {open} {isStreaming} meta={readMediaMeta} {onToggle}>
{#snippet titleSnippet()}
<span class="text-muted-foreground">Read media </span>
<span class="font-mono">{readMediaMeta?.fileName}</span>
{/snippet}
{#snippet children(_meta, _ctx)}
{#if section.toolResult}
{#if !mediaAttachment}
<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
Media attachment not found in message extras
</div>
{:else if mediaAttachment.type === AttachmentType.AUDIO}
<div class="mt-2">
<audio controls class="w-full rounded-lg">
<source
src={createBase64DataUrl(audioMimeType, mediaAttachment.base64Data)}
type={audioMimeType}
/>
Your browser does not support the audio element.
</audio>
</div>
{:else}
<div class="mt-2">
<img
src={mediaAttachment.base64Url}
alt={readMediaMeta?.fileName ?? 'media'}
class="max-h-[60vh] max-w-full rounded-lg object-contain shadow-lg"
loading="lazy"
/>
</div>
{/if}
{#if readMediaMeta?.sizeBytes || readMediaMeta?.mimeType}
<div class="mt-2 flex gap-4 text-xs text-muted-foreground">
{#if readMediaMeta?.sizeBytes}
<span>Size: {readMediaMeta.sizeBytes} bytes</span>
{/if}
{#if readMediaMeta?.mimeType}
<span>MIME: {readMediaMeta.mimeType}</span>
{/if}
</div>
{/if}
{#if readMediaMeta?.path}
<div class="mt-1 font-mono text-xs text-muted-foreground/60">{readMediaMeta.path}</div>
{/if}
{:else}
<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
Waiting for media data...
</div>
{/if}
{/snippet}
</ToolCallBlock>

56
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/read-media.ts Normal file
View file

@ -0,0 +1,56 @@
import { FILE_PATH_SEPARATOR_REGEX, NEWLINE } from '$lib/constants/code';
import {
PREFIX_FILE,
PREFIX_MIME,
PREFIX_SIZE,
READ_MEDIA_SIZE_REGEX
} from '$lib/constants/read-media';
import type { AgenticSection } from '$lib/utils';
export interface ReadMediaMeta {
fileName: string;
path: string;
sizeBytes?: number;
mimeType?: string;
}
/**
* Parse read_media tool result to extract metadata.
* Expected format (after extractBase64Attachments processing):
* File: /path/to/file.png
* Size: 12345 bytes
* MIME: image/png
* [Attachment saved: mcp-attachment-xxx.png]
*
* The data URI line is replaced by the attachment marker by
* agenticStore.extractBase64Attachments before storage.
*/
export function parseReadMediaMeta(section: AgenticSection): ReadMediaMeta | null {
if (!section.toolResult) return null;
const lines = section.toolResult.split(NEWLINE);
let fileName = '';
let path = '';
let sizeBytes: number | undefined;
let mimeType: string | undefined;
for (const line of lines) {
const trimmed = line.trim();
if (trimmed.startsWith(PREFIX_FILE)) {
path = trimmed.slice(PREFIX_FILE.length).trim();
fileName = path.split(FILE_PATH_SEPARATOR_REGEX).pop() ?? path;
} else if (trimmed.startsWith(PREFIX_SIZE)) {
const match = trimmed.match(READ_MEDIA_SIZE_REGEX);
if (match) sizeBytes = Number(match[1]);
} else if (trimmed.startsWith(PREFIX_MIME)) {
mimeType = trimmed.slice(PREFIX_MIME.length).trim();
}
}
if (!path) return null;
return { fileName, mimeType, path, sizeBytes };
}

2
tools/ui/src/lib/constants/built-in-tools.ts
View file

@ -10,6 +10,7 @@
import {
Braces,
Clock,
Eye,
FilePen,
FilePlus,
FileSearch,
@ -47,6 +48,7 @@ export const BUILTIN_TOOL_UI: Readonly<Record<BuiltInTool, BuiltinToolUiEntry>>
source: ToolSource.BUILTIN
},
[BuiltInTool.READ_FILE]: { icon: FileText, label: 'Read file', source: ToolSource.BUILTIN },
[BuiltInTool.READ_MEDIA]: { icon: Eye, label: 'Read media', source: ToolSource.FRONTEND },
[BuiltInTool.RUN_JAVASCRIPT]: {
icon: Braces,
label: 'Run JavaScript',

3
tools/ui/src/lib/constants/code.ts
View file

@ -18,6 +18,9 @@ export const TRIM_TRAILING_PADDING_REGEX = /(?:\n[ \t]*)+$/;
// `C:\foo\bar.txt`. Used wherever a parameter accepts a user-supplied path.
export const FILE_PATH_SEPARATOR_REGEX = /[\\/]/;
// Separates a file name from its extension, e.g. the '.' in `cover.png`.
export const FILE_EXTENSION_SEPARATOR = '.';
// Matches the `text:` prefix that file-type identifiers use to denote a
// plain-text language (e.g. `text:typescript`). Used by tool-call renderers
// to recover the underlying highlight.js language.

1
tools/ui/src/lib/constants/index.ts
View file

@ -49,6 +49,7 @@ export * from './sse';
export * from './precision';
export * from './processing-info';
export * from './pwa';
export * from './read-media';
export * from './routes';
export * from './sandbox';
export * from './settings-keys';

20
tools/ui/src/lib/constants/mcp-resource.ts
View file

@ -1,4 +1,4 @@
import { MimeTypeImage } from '$lib/enums';
import { MimeTypeAudio, MimeTypeImage } from '$lib/enums';
// File extension patterns for resource type detection
export const IMAGE_FILE_EXTENSION_REGEX = /\.(png|jpg|jpeg|gif|svg|webp)$/i;
@ -27,6 +27,9 @@ export const MCP_RESOURCE_ATTACHMENT_ID_PREFIX = 'res';
// Default file extension for unknown image types
export const DEFAULT_IMAGE_EXTENSION = 'img';
// Default file extension for unknown audio types
export const DEFAULT_AUDIO_EXTENSION = 'mp3';
// Default filename for resource content downloads
export const DEFAULT_RESOURCE_FILENAME = 'resource.txt';
@ -53,3 +56,18 @@ export const IMAGE_MIME_TO_EXTENSION: Record<string, string> = {
[MimeTypeImage.PNG]: 'png',
[MimeTypeImage.WEBP]: 'webp'
} as const;
/**
* Mapping from audio MIME types to file extensions.
* Used for generating attachment filenames from MIME types.
*/
export const AUDIO_MIME_TO_EXTENSION: Record<string, string> = {
[MimeTypeAudio.MP3]: 'mp3',
[MimeTypeAudio.MP3_MPEG]: 'mp3',
[MimeTypeAudio.VND_WAVE]: 'wav',
[MimeTypeAudio.WAV]: 'wav',
[MimeTypeAudio.WAVE]: 'wav',
[MimeTypeAudio.X_PN_WAV]: 'wav',
[MimeTypeAudio.X_WAV]: 'wav',
[MimeTypeAudio.X_WAVE]: 'wav'
} as const;

66
tools/ui/src/lib/constants/read-media.ts Normal file
View file

@ -0,0 +1,66 @@
import {
BuiltInTool,
JsonSchemaType,
MimeTypeAudio,
MimeTypeImage,
ToolCallType
} from '$lib/enums';
import type { OpenAIToolDefinition } from '$lib/types';
export const READ_MEDIA_TOOL_NAME = BuiltInTool.READ_MEDIA;
// header lines of the tool result, parsed back by the read_media renderer
export const PREFIX_FILE = 'File: ';
export const PREFIX_SIZE = 'Size: ';
export const PREFIX_MIME = 'MIME: ';
/** Byte count of the `Size: ` header line, e.g. `Size: 12345 bytes` -> capture group 1 is `12345`. */
export const READ_MEDIA_SIZE_REGEX = new RegExp(`^${PREFIX_SIZE}\\s*(\\d+)\\s*bytes`);
/** Image extensions the tool accepts. The server decodes images with stb_image, which has no webp or tiff. */
export const READ_MEDIA_IMAGE_MIME: Record<string, string> = {
gif: MimeTypeImage.GIF,
jpeg: MimeTypeImage.JPEG,
jpg: MimeTypeImage.JPEG,
png: MimeTypeImage.PNG
} as const;
/** Audio extensions the tool accepts. The `input_audio` API only takes wav and mp3. */
export const READ_MEDIA_AUDIO_MIME: Record<string, string> = {
mp3: MimeTypeAudio.MP3_MPEG,
wav: MimeTypeAudio.WAV
} as const;
/**
* Build the read_media tool definition for the modalities the active model has.
* At least one of the two flags must be true, otherwise the tool is not offered
* at all - a model that cannot see or hear has nothing to do with the bytes.
*/
export function buildReadMediaToolDefinition(
supportsVision: boolean,
supportsAudio: boolean
): OpenAIToolDefinition {
const kinds: string[] = [];
if (supportsVision) kinds.push(`images (${Object.keys(READ_MEDIA_IMAGE_MIME).join(', ')})`);
if (supportsAudio) kinds.push(`audio (${Object.keys(READ_MEDIA_AUDIO_MIME).join(', ')})`);
return {
function: {
description: `Read a media file and attach it to the conversation so it can be perceived directly. Supports ${kinds.join(' and ')}.`,
name: READ_MEDIA_TOOL_NAME,
parameters: {
properties: {
path: {
description: 'Path to the media file',
type: JsonSchemaType.STRING
}
},
required: ['path'],
type: JsonSchemaType.OBJECT
}
},
type: ToolCallType.FUNCTION
};
}

6
tools/ui/src/lib/constants/tools.ts
View file

@ -3,6 +3,12 @@ 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';
/** HTTP header asking the server to encode a tool's output differently, e.g. read_file returning base64. Not a tool parameter, so it stays out of the definition the model sees. */
export const X_RESP_TYPE_HEADER = 'x-resp-type';
/** `X_RESP_TYPE_HEADER` value that makes read_file return the raw bytes as base64 instead of text. */
export const RESP_TYPE_BASE64 = 'base64';
export const TOOL_GROUP_LABELS = {
[ToolSource.BUILTIN]: 'Built-in',
[ToolSource.CUSTOM]: 'JSON Schema',

1
tools/ui/src/lib/enums/files.enums.ts
View file

@ -163,6 +163,7 @@ export enum FileExtensionText {
// MIME type prefixes and includes for content detection
export enum MimeTypePrefix {
IMAGE = 'image/',
AUDIO = 'audio/',
TEXT = 'text'
}

1
tools/ui/src/lib/enums/tools.enums.ts
View file

@ -37,6 +37,7 @@ export enum GlobSearchType {
*/
export enum BuiltInTool {
READ_FILE = 'read_file',
READ_MEDIA = 'read_media',
EDIT_FILE = 'edit_file',
WRITE_FILE = 'write_file',
GET_DATETIME = 'get_datetime',

2
tools/ui/src/lib/hooks/use-context-gauge.svelte.ts
View file

@ -112,7 +112,7 @@ export function useContextGauge(): UseContextGaugeReturn {
return chatStore.getConversationModel(activeMessages() as DatabaseMessage[]);
});
const isActiveModelLoaded = $derived(
activeModelId !== null && modelsStore.isModelLoaded(activeModelId)
activeModelId !== null && (!isRouterMode() || modelsStore.isModelLoaded(activeModelId))
);
const isActiveModelLoading = $derived(
activeModelId !== null && modelsStore.isModelOperationInProgress(activeModelId)

26
tools/ui/src/lib/services/chat.service.ts
View file

@ -1,4 +1,5 @@
import { settingsStore } from '../stores/settings.svelte';
import { getAudioInputFormat } from '../utils/audio-format';
import { capImageDataURLSize } from '../utils/cap-img-size';
import {
API_CHAT,
@ -20,18 +21,12 @@ import {
import {
AttachmentType,
ContentPartType,
FileTypeAudio,
MessageRole,
MimeTypeAudio,
ReasoningFormat,
StreamConnectionState
} from '$lib/enums';
import { modelsStore } from '$lib/stores/models.svelte';
import type {
AudioInputFormat,
DatabaseMessageExtraMcpPrompt,
DatabaseMessageExtraMcpResource
} from '$lib/types';
import type { DatabaseMessageExtraMcpPrompt, DatabaseMessageExtraMcpResource } from '$lib/types';
import type {
ApiChatCompletionToolCall,
ApiChatMessageContentPart,
@ -43,23 +38,6 @@ import { getAuthHeaders, getJsonHeaders } from '$lib/utils/api-headers';
import { formatAttachmentText } from '$lib/utils/formatters';
import { streamIdentity } from '$lib/utils/stream-identity';
function getAudioInputFormat(mimeType: string): AudioInputFormat {
const normalizedMimeType = mimeType.trim().toLowerCase();
if (
normalizedMimeType === MimeTypeAudio.WAV ||
normalizedMimeType === MimeTypeAudio.WAVE ||
normalizedMimeType === MimeTypeAudio.X_WAV ||
normalizedMimeType === MimeTypeAudio.X_WAVE ||
normalizedMimeType === MimeTypeAudio.VND_WAVE ||
normalizedMimeType === MimeTypeAudio.X_PN_WAV
) {
return FileTypeAudio.WAV;
}
return FileTypeAudio.MP3;
}
interface ResumableStreamState {
bytesReceived: number;
updatedAt: number;

112
tools/ui/src/lib/services/read-media.service.ts Normal file
View file

@ -0,0 +1,112 @@
import { ToolsService } from './tools.service';
import {
FILE_EXTENSION_SEPARATOR,
FILE_PATH_SEPARATOR_REGEX,
NEWLINE,
PREFIX_FILE,
PREFIX_MIME,
PREFIX_SIZE,
READ_MEDIA_AUDIO_MIME,
READ_MEDIA_IMAGE_MIME,
RESP_TYPE_BASE64
} from '$lib/constants';
import { BuiltInTool, ToolResponseField } from '$lib/enums';
import type { ToolExecutionResult } from '$lib/types';
/** Modalities of the model the tool call runs for. */
export interface ReadMediaCapabilities {
audio: boolean;
vision: boolean;
}
/** Lowercase extension of a path, without the dot. Empty when the file name has none. */
function fileExtension(path: string): string {
const name = path.split(FILE_PATH_SEPARATOR_REGEX).pop() ?? '';
const dot = name.lastIndexOf(FILE_EXTENSION_SEPARATOR);
return dot > 0 ? name.slice(dot + 1).toLowerCase() : '';
}
/**
* **ReadMediaService** - frontend executor for the `read_media` tool
*
* The tool is synthetic: no such tool exists on the server. It reads the file
* through the built-in `read_file` tool with the `base64` response type, then
* turns the bytes into a data URI line. The agentic store lifts that line into
* an image or audio attachment on the tool result message, which is what makes
* the model perceive the file instead of reading a wall of base64.
*
* Living in the frontend is what lets it exist only for models that can
* actually use the result - the server has no idea which model is selected.
*
* @see buildReadMediaToolDefinition in constants/read-media.ts - tool schema sent to the LLM
* @see agenticStore in stores/agentic.svelte.ts - tool dispatch and attachment extraction
*/
export class ReadMediaService {
static async executeTool(
params: Record<string, unknown>,
capabilities: ReadMediaCapabilities,
signal?: AbortSignal,
cwd?: string
): Promise<ToolExecutionResult> {
const path = typeof params.path === 'string' ? params.path : '';
if (!path) {
return { content: 'Error: missing "path" argument.', isError: true };
}
const extension = fileExtension(path);
const imageMime = READ_MEDIA_IMAGE_MIME[extension];
const audioMime = READ_MEDIA_AUDIO_MIME[extension];
let resolvedMime: string | undefined;
if (imageMime && capabilities.vision) resolvedMime = imageMime;
else if (audioMime && capabilities.audio) resolvedMime = audioMime;
if (!resolvedMime) {
const supported = [
...(capabilities.vision ? Object.keys(READ_MEDIA_IMAGE_MIME) : []),
...(capabilities.audio ? Object.keys(READ_MEDIA_AUDIO_MIME) : [])
];
// an unreadable-by-this-model file is a dead end, so say why instead of failing silently
const reason =
imageMime || audioMime
? `the current model cannot perceive ".${extension}" files`
: `".${extension}" is not a supported media type`;
return {
content: `Error: ${reason}. Supported: ${supported.join(', ')}.`,
isError: true
};
}
const raw = await ToolsService.executeToolRaw(
BuiltInTool.READ_FILE,
{ path },
signal,
cwd,
RESP_TYPE_BASE64
);
if (ToolResponseField.ERROR in raw) {
return { content: String(raw[ToolResponseField.ERROR]), isError: true };
}
const base64 = typeof raw.base64 === 'string' ? raw.base64 : '';
if (!base64) {
return { content: `Error: no data returned for ${path}.`, isError: true };
}
const sizeBytes = typeof raw.size_bytes === 'number' ? raw.size_bytes : 0;
const content = [
`${PREFIX_FILE}${path}`,
`${PREFIX_SIZE}${sizeBytes} bytes`,
`${PREFIX_MIME}${resolvedMime}`,
`data:${resolvedMime};base64,${base64}`
].join(NEWLINE);
return { content, isError: false };
}
}

16
tools/ui/src/lib/services/tools.service.ts
View file

@ -1,5 +1,5 @@
import { base } from '$app/paths';
import { API_TOOLS, X_TOOL_CWD_HEADER } from '$lib/constants';
import { API_TOOLS, X_RESP_TYPE_HEADER, X_TOOL_CWD_HEADER } from '$lib/constants';
import { ToolResponseField } from '$lib/enums';
import type { ServerBuiltinToolInfo, ToolExecutionResult } from '$lib/types';
import { apiFetch } from '$lib/utils';
@ -51,16 +51,26 @@ export class ToolsService {
* 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.
*
* @param respType - sent as the x-resp-type request header. Only read_file
* honors it, with `base64` to get the raw bytes instead of decoded text.
*/
static async executeToolRaw(
toolName: string,
params: Record<string, unknown>,
signal?: AbortSignal,
cwd?: string
cwd?: string,
respType?: string
): Promise<Record<string, unknown>> {
const headers: Record<string, string> = {};
if (cwd) headers[X_TOOL_CWD_HEADER] = cwd;
if (respType) headers[X_RESP_TYPE_HEADER] = respType;
return apiFetch<Record<string, unknown>>(API_TOOLS.EXECUTE, {
body: JSON.stringify({ params, tool: toolName }),
headers: cwd ? { [X_TOOL_CWD_HEADER]: cwd } : undefined,
headers: Object.keys(headers).length > 0 ? headers : undefined,
method: 'POST',
signal
});

49
tools/ui/src/lib/stores/agentic.svelte.ts
View file

@ -22,7 +22,9 @@
import { DEFAULT_AGENTIC_CONFIG, NEWLINE } from '$lib/constants';
import {
AUDIO_MIME_TO_EXTENSION,
DATA_URI_BASE64_REGEX,
DEFAULT_AUDIO_EXTENSION,
DEFAULT_IMAGE_EXTENSION,
IMAGE_MIME_TO_EXTENSION,
MCP_ATTACHMENT_NAME_PREFIX
@ -36,6 +38,7 @@ import {
ToolCallType
} from '$lib/enums';
import { ChatService } from '$lib/services';
import { ReadMediaService } from '$lib/services/read-media.service';
import { SandboxService } from '$lib/services/sandbox.service';
import { ToolsService } from '$lib/services/tools.service';
import { conversationsStore } from '$lib/stores/conversations.svelte';
@ -75,9 +78,10 @@ import type {
import type {
DatabaseMessage,
DatabaseMessageExtra,
DatabaseMessageExtraAudioFile,
DatabaseMessageExtraImageFile
} from '$lib/types/database';
import { isAbortError } from '$lib/utils';
import { getAudioInputFormat, isAbortError } from '$lib/utils';
import { SvelteMap } from 'svelte/reactivity';
function createDefaultSession(): AgenticSession {
@ -900,7 +904,18 @@ class AgenticStore {
if (executionResult.isError) toolSuccess = false;
} else if (toolSource === ToolSource.FRONTEND) {
const args = this.parseToolArguments(toolCall.function.arguments);
const executionResult = await SandboxService.executeTool(toolName, args, signal);
const executionResult =
toolName === BuiltInTool.READ_MEDIA
? await ReadMediaService.executeTool(
args,
{
audio: modelsStore.modelSupportsAudio(effectiveModel),
vision: modelsStore.modelSupportsVision(effectiveModel)
},
signal,
conversationsStore.activeConversation?.cwd
)
: await SandboxService.executeTool(toolName, args, signal);
result = executionResult.content;
@ -990,7 +1005,19 @@ class AgenticStore {
];
for (const attachment of attachments) {
if (attachment.type === AttachmentType.IMAGE) {
if (attachment.type === AttachmentType.AUDIO) {
if (modelsStore.modelSupportsAudio(effectiveModel)) {
contentParts.push({
input_audio: {
data: (attachment as DatabaseMessageExtraAudioFile).base64Data,
format: getAudioInputFormat(
(attachment as DatabaseMessageExtraAudioFile).mimeType
)
},
type: ContentPartType.INPUT_AUDIO
});
}
} else if (attachment.type === AttachmentType.IMAGE) {
if (modelsStore.modelSupportsVision(effectiveModel)) {
contentParts.push({
image_url: {
@ -1101,6 +1128,18 @@ class AgenticStore {
return `[Attachment saved: ${name}]`;
}
if (mimeType.startsWith(MimeTypePrefix.AUDIO)) {
// audio extras hold the bare base64, the input_audio part has no room for a data URI
attachments.push({
base64Data,
mimeType,
name,
type: AttachmentType.AUDIO
});
return `[Attachment saved: ${name}]`;
}
return line;
});
@ -1108,7 +1147,9 @@ class AgenticStore {
}
private buildAttachmentName(mimeType: string, index: number): string {
const extension = IMAGE_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_IMAGE_EXTENSION;
const extension = mimeType.startsWith(MimeTypePrefix.AUDIO)
? (AUDIO_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_AUDIO_EXTENSION)
: (IMAGE_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_IMAGE_EXTENSION);
return `${MCP_ATTACHMENT_NAME_PREFIX}-${Date.now()}-${index}.${extension}`;
}

2
tools/ui/src/lib/stores/context-gauge-popup.svelte.ts
View file

@ -16,7 +16,7 @@ import {
let closeTimer: ReturnType<typeof setTimeout> | undefined;
let lastPointerType = '';
export const gaugePopup = $state({ bottom: 0, centerX: 0, open: false });
export const gaugePopup = $state({ bottom: 0, centerX: 0, detailsOpen: false, open: false });
function openFrom(trigger: HTMLElement): void {
clearTimeout(closeTimer);

5
tools/ui/src/lib/stores/models.svelte.ts
View file

@ -18,7 +18,10 @@ import { ModelsService } from '$lib/services/models.service';
import { PropsService } from '$lib/services/props.service';
import { conversationsStore } from '$lib/stores/conversations.svelte';
import { isRouterMode, serverStore } from '$lib/stores/server.svelte';
import { getAuthHeaders, TTLCache } from '$lib/utils';
// deep imports, not the '$lib/utils' barrel: it re-exports modules that reach back
// into the stores, and going through it here would read a half-built module
import { getAuthHeaders } from '$lib/utils/api-headers';
import { TTLCache } from '$lib/utils/cache-ttl';
import {
detectThinkingSupport,
detectThinkingSupportWithReason

41
tools/ui/src/lib/stores/tools.svelte.ts
View file

@ -1,4 +1,5 @@
import {
buildReadMediaToolDefinition,
buildSandboxToolDefinition,
DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY,
HOME_TILDE,
@ -15,6 +16,7 @@ import {
} from '$lib/enums';
import { ToolsService } from '$lib/services/tools.service';
import { mcpStore } from '$lib/stores/mcp.svelte';
import { modelsStore, selectedModelName } from '$lib/stores/models.svelte';
import { config } from '$lib/stores/settings.svelte';
import type { OpenAIToolDefinition, ToolEntry, ToolGroup } from '$lib/types';
import { SvelteMap, SvelteSet } from 'svelte/reactivity';
@ -168,9 +170,42 @@ class ToolsStore {
}
get frontendTools(): OpenAIToolDefinition[] {
return config().jsSandboxEnabled
? [buildSandboxToolDefinition(!!config().symbolicMathEnabled)]
: [];
const tools: OpenAIToolDefinition[] = [];
if (config().jsSandboxEnabled) {
tools.push(buildSandboxToolDefinition(!!config().symbolicMathEnabled));
}
const readMedia = this.readMediaTool();
if (readMedia) tools.push(readMedia);
return tools;
}
/**
* `read_media` runs in the frontend on top of the server's `read_file`, so it
* exists only when that tool is served and the active model can perceive the
* bytes. The server cannot make this call - it does not know which model the
* conversation uses.
*/
private readMediaTool(): OpenAIToolDefinition | null {
const hasReadFile = this._builtinTools.some(
(def) => def.function.name === BuiltInTool.READ_FILE
);
if (!hasReadFile) return null;
const model = selectedModelName() ?? modelsStore.models[0]?.model ?? '';
if (!model) return null;
const vision = modelsStore.modelSupportsVision(model);
const audio = modelsStore.modelSupportsAudio(model);
if (!vision && !audio) return null;
return buildReadMediaToolDefinition(vision, audio);
}
get customTools(): OpenAIToolDefinition[] {

19
tools/ui/src/lib/utils/agentic.ts
View file

@ -50,11 +50,11 @@ export interface AgenticSection {
}
/**
* Represents a tool result line that may reference an image attachment
* Represents a tool result line that may reference a media attachment (image or audio)
*/
export type ToolResultLine = {
text: string;
image?: DatabaseMessageExtraImageFile;
media?: DatabaseMessageExtraImageFile | DatabaseMessageExtraAudioFile;
};
/**
@ -301,16 +301,16 @@ export function splitSearchSummaryList(
return { lines };
}
/** Bounded cache for parseToolResultWithImages results. */
/** Bounded cache for parseToolResultWithMedia results. */
const TOOL_RESULT_LINES_CACHE_MAX_SIZE = 32;
const toolResultLinesCache = new Map<string, ToolResultLine[]>();
/**
* Parse tool result text into lines, matching image attachments by name.
* Parse tool result text into lines, matching media attachments (images and audio) by name.
* Memoized: called per render during streaming on unchanged tool result
* strings with unchanged extras.
*/
export function parseToolResultWithImages(
export function parseToolResultWithMedia(
toolResult: string,
extras?: DatabaseMessageExtra[]
): ToolResultLine[] {
@ -332,12 +332,13 @@ export function parseToolResultWithImages(
if (!match || !extras) return { text: line };
const attachmentName = match[1];
const image = extras.find(
(e): e is DatabaseMessageExtraImageFile =>
e.type === AttachmentType.IMAGE && e.name === attachmentName
const media = extras.find(
(e): e is DatabaseMessageExtraImageFile | DatabaseMessageExtraAudioFile =>
(e.type === AttachmentType.IMAGE || e.type === AttachmentType.AUDIO) &&
e.name === attachmentName
);
return { image, text: line };
return { media, text: line };
});
if (toolResultLinesCache.size >= TOOL_RESULT_LINES_CACHE_MAX_SIZE) {

22
tools/ui/src/lib/utils/audio-format.ts Normal file
View file

@ -0,0 +1,22 @@
import { FileTypeAudio, MimeTypeAudio } from '$lib/enums';
import type { AudioInputFormat } from '$lib/types/api';
/**
* Map a MIME type to the AudioInputFormat expected by the API.
*/
export function getAudioInputFormat(mimeType: string): AudioInputFormat {
const normalizedMimeType = mimeType.trim().toLowerCase();
if (
normalizedMimeType === MimeTypeAudio.WAV ||
normalizedMimeType === MimeTypeAudio.WAVE ||
normalizedMimeType === MimeTypeAudio.X_WAV ||
normalizedMimeType === MimeTypeAudio.X_WAVE ||
normalizedMimeType === MimeTypeAudio.VND_WAVE ||
normalizedMimeType === MimeTypeAudio.X_PN_WAV
) {
return FileTypeAudio.WAV;
}
return FileTypeAudio.MP3;
}

5
tools/ui/src/lib/utils/index.ts
View file

@ -248,7 +248,7 @@ export {
export {
deriveAgenticSections,
buildAssistantRawOutput,
parseToolResultWithImages,
parseToolResultWithMedia,
splitSearchSummaryList,
hasAgenticContent,
classifyToolResult,
@ -325,3 +325,6 @@ export { uuid } from './uuid';
// CSS utilities
export { remToPx } from './css';
// Audio format helper (used by agentic store and chat service)
export { getAudioInputFormat } from './audio-format';

12
tools/ui/tests/client/README-perf.md
View file

@ -33,12 +33,12 @@ npx vitest --project=client --run tests/client/agentic-stream.perf.svelte.test.t
The point of the harness is the _scaling curve_, not any single number.
| Knob | Reads on |
| --------------------------- | ---------------------------------------------------------------------------------------------------- |
| `priorToolCalls` (0/1/5/20) | the reactive fan-out. Flat => no fan-out. Linear => confirmed. |
| `toolResultBytes` | whole-blob string scans (`extractSearchResults`, `parseToolResultWithImages`, `classifyToolResult`). |
| `editFileEdits` | `computeLineDiff`, the O(m\*n) LCS. |
| `openCodeFence` | `hljs.highlightAuto` on partial code. |
| Knob | Reads on |
| --------------------------- | --------------------------------------------------------------------------------------------------- |
| `priorToolCalls` (0/1/5/20) | the reactive fan-out. Flat => no fan-out. Linear => confirmed. |
| `toolResultBytes` | whole-blob string scans (`extractSearchResults`, `parseToolResultWithMedia`, `classifyToolResult`). |
| `editFileEdits` | `computeLineDiff`, the O(m\*n) LCS. |
| `openCodeFence` | `hljs.highlightAuto` on partial code. |
Deliberately no hard assertions: CI timing is noisy and the value here is the
before/after delta, not a gate.

10
tools/ui/tests/unit/agentic-hotpath.bench.ts
View file

@ -6,7 +6,7 @@
//
// Run: npx vitest bench --project=unit tests/unit/agentic-hotpath.bench.ts
import { classifyToolResult, parseToolResultWithImages } from '$lib/utils/agentic';
import { classifyToolResult, parseToolResultWithMedia } from '$lib/utils/agentic';
import { detectIncompleteCodeBlock, highlightCode } from '$lib/utils/code';
import { computeLineDiff } from '$lib/utils/compute-line-diff';
import { preprocessLaTeX } from '$lib/utils/latex-protection';
@ -200,17 +200,17 @@ describe('exit-code regex', () => {
// --- per-line result parsers ----------------------------------------------
describe('parseToolResultWithImages', () => {
describe('parseToolResultWithMedia', () => {
bench('1KB', () => {
parseToolResultWithImages(SHELL_OUTPUT_1KB, []);
parseToolResultWithMedia(SHELL_OUTPUT_1KB, []);
});
bench('200KB', () => {
parseToolResultWithImages(SHELL_OUTPUT_200KB, []);
parseToolResultWithMedia(SHELL_OUTPUT_200KB, []);
});
bench('2MB', () => {
parseToolResultWithImages(SHELL_OUTPUT_2MB, []);
parseToolResultWithMedia(SHELL_OUTPUT_2MB, []);
});
});