Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-09-20 01:31:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

qwen3tts support 1.7b model

Browse Source
This commit is contained in:
Concedo
2026-03-13 21:15:24 +08:00
parent a13641c00c
commit 4189508ef3
3 changed files with 356 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files
otherarch/qwen3tts/q3tts_convert_tts_to_gguf.py
+43 -32
View File
@@ -1,6 +1,6 @@
#!/usr/bin/env python3
"""
Convert HuggingFace Qwen3-TTS-12Hz-0.6B-Base model to GGUF format.
Convert HuggingFace Qwen3-TTS model to GGUF format.
Usage:
python scripts/convert_tts_to_gguf.py \
@@ -39,7 +39,7 @@ logger = logging.getLogger(__name__)
class Qwen3TTSConverter:
"""Converter for Qwen3-TTS-12Hz-0.6B-Base model to GGUF format."""
"""Converter for Qwen3-TTS models to GGUF format."""
# Direct tensor name mapping from HuggingFace to GGML conventions
TENSOR_MAP = {
@@ -55,6 +55,9 @@ class Qwen3TTSConverter:
"talker.text_projection.linear_fc2.bias": "talker.text_proj.fc2.bias",
# Code Predictor - Output norm
"talker.code_predictor.model.norm.weight": "code_pred.output_norm.weight",
# Code Predictor - MTP projection (talker_hidden -> code_pred_hidden)
"talker.code_predictor.small_to_mtp_projection.weight": "code_pred.mtp_proj.weight",
"talker.code_predictor.small_to_mtp_projection.bias": "code_pred.mtp_proj.bias",
# Speaker Encoder - Initial conv
"speaker_encoder.blocks.0.conv.weight": "spk_enc.conv0.weight",
"speaker_encoder.blocks.0.conv.bias": "spk_enc.conv0.bias",
@@ -177,6 +180,11 @@ class Qwen3TTSConverter:
# Code Predictor parameters
self.code_predictor_num_layers = code_predictor_config.get("num_hidden_layers", 5)
self.code_predictor_vocab_size = code_predictor_config.get("vocab_size", 2048)
self.code_pred_hidden_size = code_predictor_config.get("hidden_size", 1024)
self.code_pred_n_attention_heads = code_predictor_config.get("num_attention_heads", 16)
self.code_pred_n_key_value_heads = code_predictor_config.get("num_key_value_heads", 8)
self.code_pred_intermediate_size = code_predictor_config.get("intermediate_size", 3072)
self.code_pred_head_dim = code_predictor_config.get("head_dim", 128)
# Speaker Encoder parameters
self.speaker_enc_dim = speaker_encoder_config.get("enc_dim", 1024)
@@ -187,8 +195,10 @@ class Qwen3TTSConverter:
self.codec_bos_id = talker_config.get("codec_bos_id", 2149)
self.codec_eos_id = talker_config.get("codec_eos_token_id", 2150)
# Model name
self.model_name = "Qwen3-TTS-12Hz-0.6B"
# Model name — derive from config or input directory name
self.model_name = self.config.get("general_name",
self.input_dir.name.replace("-Base", "").replace("-", " ").strip()
or "Qwen3-TTS")
def _map_tensor_name(self, hf_name: str) -> str | None:
"""Map HuggingFace tensor name to GGML convention."""
@@ -241,31 +251,41 @@ class Qwen3TTSConverter:
for name in f.keys():
yield name, f.get_tensor(name)
def _should_quantize(self, tensor_name: str) -> bool:
def _should_quantize(self, tensor_name: str, n_dims: int) -> bool:
"""Determine if a tensor should be quantized (Q8_0) or kept in F16.
Tensors to keep in F16 for quality:
Tensors to keep in F16:
- 3D tensors (conv1d weights — Q8_0 doesn't support 3D)
- Speaker encoder tensors (tiny, not worth quantizing)
- Embeddings (codec_embd, text_embd, codebook)
- Layer norms (attn_norm, ffn_norm, output_norm)
- Biases
- LM heads
"""
# 3D tensors (conv1d) can't be quantized
if n_dims >= 3:
return False
# Speaker encoder is tiny, keep in F16
if "spk_enc" in tensor_name:
return False
# Keep embeddings in F16
if any(x in tensor_name for x in ["_embd", "codebook"]):
return False
# Keep layer norms in F16
if "_norm" in tensor_name:
return False
# Keep biases in F16
if ".bias" in tensor_name:
return False
# Keep LM heads in F16
if "lm_head" in tensor_name or "codec_head" in tensor_name:
return False
# Quantize weight matrices
return True
@@ -292,10 +312,10 @@ class Qwen3TTSConverter:
elif self.output_type == "f16":
return data.astype(np.float16), gguf.GGMLQuantizationType.F16
elif self.output_type == "q8_0":
if not self._should_quantize(tensor_name):
if not self._should_quantize(tensor_name, n_dims):
logger.debug(f"Keeping {tensor_name} in F16 (not quantizing)")
return data.astype(np.float16), gguf.GGMLQuantizationType.F16
data = data.astype(np.float32)
try:
quantized = gguf.quants.quantize(data, gguf.GGMLQuantizationType.Q8_0)
@@ -303,18 +323,6 @@ class Qwen3TTSConverter:
except Exception as e:
logger.warning(f"Q8_0 quantization failed for {tensor_name}: {e}, falling back to F16")
return data.astype(np.float16), gguf.GGMLQuantizationType.F16
elif self.output_type == "q4_k":
if not self._should_quantize(tensor_name):
logger.debug(f"Keeping {tensor_name} in F16 (not quantizing)")
return data.astype(np.float16), gguf.GGMLQuantizationType.F16
data = data.astype(np.float32)
try:
quantized = gguf.quants.quantize(data, gguf.GGMLQuantizationType.Q4_K)
return quantized, gguf.GGMLQuantizationType.Q4_K
except Exception as e:
logger.warning(f"Q4_K quantization failed for {tensor_name}: {e}, falling back to F16")
return data.astype(np.float16), gguf.GGMLQuantizationType.F16
else:
return data.astype(np.float16), gguf.GGMLQuantizationType.F16
@@ -416,7 +424,7 @@ class Qwen3TTSConverter:
def _add_metadata(self, writer: gguf.GGUFWriter) -> None:
"""Add model metadata to GGUF writer."""
arch = "qwen3-tts"
# General metadata
writer.add_name(self.model_name)
writer.add_type(gguf.GGUFType.MODEL)
@@ -428,8 +436,6 @@ class Qwen3TTSConverter:
ftype = gguf.LlamaFileType.MOSTLY_F16
elif self.output_type == "q8_0":
ftype = gguf.LlamaFileType.MOSTLY_Q8_0
elif self.output_type == "q4_k":
ftype = gguf.LlamaFileType.MOSTLY_Q4_K_M
else:
ftype = gguf.LlamaFileType.MOSTLY_F16
writer.add_file_type(ftype)
@@ -460,6 +466,11 @@ class Qwen3TTSConverter:
# Code Predictor parameters
writer.add_uint32(f"{arch}.code_predictor.layer_count", self.code_predictor_num_layers)
writer.add_uint32(f"{arch}.code_predictor.vocab_size", self.code_predictor_vocab_size)
writer.add_uint32(f"{arch}.code_pred.embedding_length", self.code_pred_hidden_size)
writer.add_uint32(f"{arch}.code_pred.attention.head_count", self.code_pred_n_attention_heads)
writer.add_uint32(f"{arch}.code_pred.attention.head_count_kv", self.code_pred_n_key_value_heads)
writer.add_uint32(f"{arch}.code_pred.feed_forward_length", self.code_pred_intermediate_size)
writer.add_uint32(f"{arch}.code_pred.attention.key_length", self.code_pred_head_dim)
# Speaker Encoder parameters
writer.add_uint32(f"{arch}.speaker_encoder.embedding_length", self.speaker_enc_dim)
@@ -526,7 +537,7 @@ class Qwen3TTSConverter:
def main():
parser = argparse.ArgumentParser(
description="Convert Qwen3-TTS-12Hz-0.6B-Base model to GGUF format"
description="Convert Qwen3-TTS model to GGUF format"
)
parser.add_argument(
"--input", "-i",
@@ -542,9 +553,9 @@ def main():
)
parser.add_argument(
"--type", "-t",
choices=["f16", "f32", "q8_0", "q4_k"],
choices=["f16", "f32", "q8_0"],
default="f16",
help="Output data type (default: f16). q8_0 provides ~50%% size reduction, q4_k provides ~70%% size reduction."
help="Output data type (default: f16). q8_0 provides ~50%% size reduction."
)
parser.add_argument(
"--verbose", "-v",
@@ -566,4 +577,4 @@ def main():
if __name__ == "__main__":
main()
main()
otherarch/qwen3tts/tts_transformer.cpp
+278 -128
View File
@@ -23,6 +23,16 @@ namespace qwen3_tts {
TTSTransformer::TTSTransformer() = default;
struct ggml_tensor * TTSTransformer::mul_mat(struct ggml_context * ctx,
struct ggml_tensor * a,
struct ggml_tensor * b) {
struct ggml_tensor * result = ggml_mul_mat(ctx, a, b);
if (force_f32_acc_) {
ggml_mul_mat_set_prec(result, GGML_PREC_F32);
}
return result;
}
TTSTransformer::~TTSTransformer() {
unload_model();
}
@@ -227,6 +237,22 @@ bool TTSTransformer::parse_config(struct gguf_context * ctx) {
"qwen3-tts.code_predictor.vocab_size",
}, 2048);
cfg.code_pred_hidden_size = get_u32_any({
"qwen3-tts.code_pred.embedding_length",
}, cfg.hidden_size); // default to talker hidden_size (0.6B case)
cfg.code_pred_n_attention_heads = get_u32_any({
"qwen3-tts.code_pred.attention.head_count",
}, cfg.n_attention_heads);
cfg.code_pred_n_key_value_heads = get_u32_any({
"qwen3-tts.code_pred.attention.head_count_kv",
}, cfg.n_key_value_heads);
cfg.code_pred_intermediate_size = get_u32_any({
"qwen3-tts.code_pred.feed_forward_length",
}, cfg.intermediate_size);
cfg.code_pred_head_dim = get_u32_any({
"qwen3-tts.code_pred.attention.key_length",
}, cfg.head_dim);
cfg.codec_pad_id = get_u32_any({
"qwen3-tts.codec.pad_id",
}, 2148);
@@ -402,44 +428,44 @@ bool TTSTransformer::create_tensors(struct gguf_context * ctx) {
layer_idx >= 0 && layer_idx < cfg.code_pred_layers) {
if (strstr(name, "attn_norm.weight")) {
ne[0] = cfg.hidden_size;
ne[0] = cfg.code_pred_hidden_size;
n_dims = 1;
} else if (strstr(name, "attn_q_norm.weight")) {
ne[0] = cfg.head_dim;
ne[0] = cfg.code_pred_head_dim;
n_dims = 1;
} else if (strstr(name, "attn_k_norm.weight")) {
ne[0] = cfg.head_dim;
ne[0] = cfg.code_pred_head_dim;
n_dims = 1;
} else if (strstr(name, "attn_q.weight")) {
ne[0] = cfg.hidden_size;
ne[1] = cfg.n_attention_heads * cfg.head_dim;
ne[0] = cfg.code_pred_hidden_size;
ne[1] = cfg.code_pred_n_attention_heads * cfg.code_pred_head_dim;
n_dims = 2;
} else if (strstr(name, "attn_k.weight")) {
ne[0] = cfg.hidden_size;
ne[1] = cfg.n_key_value_heads * cfg.head_dim;
ne[0] = cfg.code_pred_hidden_size;
ne[1] = cfg.code_pred_n_key_value_heads * cfg.code_pred_head_dim;
n_dims = 2;
} else if (strstr(name, "attn_v.weight")) {
ne[0] = cfg.hidden_size;
ne[1] = cfg.n_key_value_heads * cfg.head_dim;
ne[0] = cfg.code_pred_hidden_size;
ne[1] = cfg.code_pred_n_key_value_heads * cfg.code_pred_head_dim;
n_dims = 2;
} else if (strstr(name, "attn_output.weight")) {
ne[0] = cfg.n_attention_heads * cfg.head_dim;
ne[1] = cfg.hidden_size;
ne[0] = cfg.code_pred_n_attention_heads * cfg.code_pred_head_dim;
ne[1] = cfg.code_pred_hidden_size;
n_dims = 2;
} else if (strstr(name, "ffn_norm.weight")) {
ne[0] = cfg.hidden_size;
ne[0] = cfg.code_pred_hidden_size;
n_dims = 1;
} else if (strstr(name, "ffn_gate.weight")) {
ne[0] = cfg.hidden_size;
ne[1] = cfg.intermediate_size;
ne[0] = cfg.code_pred_hidden_size;
ne[1] = cfg.code_pred_intermediate_size;
n_dims = 2;
} else if (strstr(name, "ffn_up.weight")) {
ne[0] = cfg.hidden_size;
ne[1] = cfg.intermediate_size;
ne[0] = cfg.code_pred_hidden_size;
ne[1] = cfg.code_pred_intermediate_size;
n_dims = 2;
} else if (strstr(name, "ffn_down.weight")) {
ne[0] = cfg.intermediate_size;
ne[1] = cfg.hidden_size;
ne[0] = cfg.code_pred_intermediate_size;
ne[1] = cfg.code_pred_hidden_size;
n_dims = 2;
} else {
continue;
@@ -447,6 +473,17 @@ bool TTSTransformer::create_tensors(struct gguf_context * ctx) {
} else {
continue;
}
} else if (strstr(name, "code_pred.mtp_proj.")) {
if (strstr(name, "weight")) {
ne[0] = cfg.hidden_size;
ne[1] = cfg.code_pred_hidden_size;
n_dims = 2;
} else if (strstr(name, "bias")) {
ne[0] = cfg.code_pred_hidden_size;
n_dims = 1;
} else {
continue;
}
} else if (strstr(name, "code_pred.codec_embd.")) {
int cb_idx = -1;
if (sscanf(name, "code_pred.codec_embd.%d.weight", &cb_idx) == 1 &&
@@ -464,7 +501,7 @@ bool TTSTransformer::create_tensors(struct gguf_context * ctx) {
int cb_idx = -1;
if (sscanf(name, "code_pred.lm_head.%d.weight", &cb_idx) == 1 &&
cb_idx >= 0 && cb_idx < cfg.n_codebooks - 1) {
ne[0] = cfg.hidden_size;
ne[0] = cfg.code_pred_hidden_size;
ne[1] = cfg.code_pred_vocab_size;
n_dims = 2;
} else {
@@ -474,7 +511,7 @@ bool TTSTransformer::create_tensors(struct gguf_context * ctx) {
if (skip_ggml_code_pred_layers_) {
continue;
}
ne[0] = cfg.hidden_size;
ne[0] = cfg.code_pred_hidden_size;
n_dims = 1;
} else {
continue;
@@ -552,6 +589,10 @@ bool TTSTransformer::create_tensors(struct gguf_context * ctx) {
}
} else if (strstr(name, "code_pred.output_norm.weight")) {
model_.code_pred_output_norm = tensor;
} else if (strstr(name, "code_pred.mtp_proj.weight")) {
model_.code_pred_mtp_proj_w = tensor;
} else if (strstr(name, "code_pred.mtp_proj.bias")) {
model_.code_pred_mtp_proj_b = tensor;
}
}
@@ -680,8 +721,8 @@ bool TTSTransformer::init_code_pred_kv_cache(int32_t n_ctx) {
state_.code_pred_cache.n_ctx = n_ctx;
state_.code_pred_cache.n_used = 0;
state_.code_pred_cache.head_dim = cfg.head_dim;
state_.code_pred_cache.n_kv_heads = cfg.n_key_value_heads;
state_.code_pred_cache.head_dim = cfg.code_pred_head_dim;
state_.code_pred_cache.n_kv_heads = cfg.code_pred_n_key_value_heads;
state_.code_pred_cache.n_layers = cfg.code_pred_layers;
const size_t n_tensors = cfg.code_pred_layers * 2;
@@ -705,12 +746,12 @@ bool TTSTransformer::init_code_pred_kv_cache(int32_t n_ctx) {
for (int il = 0; il < cfg.code_pred_layers; ++il) {
state_.code_pred_cache.k_cache[il] = ggml_new_tensor_3d(
state_.code_pred_cache.ctx, GGML_TYPE_F16,
cfg.head_dim, cfg.n_key_value_heads, n_ctx);
cfg.code_pred_head_dim, cfg.code_pred_n_key_value_heads, n_ctx);
ggml_format_name(state_.code_pred_cache.k_cache[il], "code_pred_k_cache_%d", il);
state_.code_pred_cache.v_cache[il] = ggml_new_tensor_3d(
state_.code_pred_cache.ctx, GGML_TYPE_F16,
cfg.head_dim, cfg.n_key_value_heads, n_ctx);
cfg.code_pred_head_dim, cfg.code_pred_n_key_value_heads, n_ctx);
ggml_format_name(state_.code_pred_cache.v_cache[il], "code_pred_v_cache_%d", il);
}
@@ -783,6 +824,11 @@ bool TTSTransformer::lookup_embedding_rows(struct ggml_tensor * embedding, const
}
struct ggml_tensor * inp = ggml_graph_get_tensor(gf, input_name);
if (!inp) {
error_msg_ = std::string("Failed to find input tensor: ") + input_name;
ggml_backend_sched_reset(state_.sched);
return false;
}
ggml_backend_tensor_set(inp, token_ids, 0, n_tokens * sizeof(int32_t));
if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) {
@@ -876,10 +922,10 @@ bool TTSTransformer::project_text_tokens(const int32_t * text_tokens, int32_t n_
ggml_set_input(inp_tokens);
struct ggml_tensor * cur = ggml_get_rows(ctx0, model_.text_embd, inp_tokens);
cur = ggml_mul_mat(ctx0, model_.text_proj_fc1, cur);
cur = mul_mat(ctx0, model_.text_proj_fc1, cur);
cur = ggml_add(ctx0, cur, model_.text_proj_fc1_bias);
cur = ggml_silu(ctx0, cur);
cur = ggml_mul_mat(ctx0, model_.text_proj_fc2, cur);
cur = mul_mat(ctx0, model_.text_proj_fc2, cur);
cur = ggml_add(ctx0, cur, model_.text_proj_fc2_bias);
ggml_set_name(cur, "text_proj_out");
@@ -893,6 +939,11 @@ bool TTSTransformer::project_text_tokens(const int32_t * text_tokens, int32_t n_
}
struct ggml_tensor * inp = ggml_graph_get_tensor(gf, "inp_text_tokens");
if (!inp) {
error_msg_ = "Failed to find inp_text_tokens tensor in graph";
ggml_backend_sched_reset(state_.sched);
return false;
}
ggml_backend_tensor_set(inp, text_tokens, 0, n_tokens * sizeof(int32_t));
if (ggml_backend_sched_graph_compute(state_.sched, gf) != GGML_STATUS_SUCCESS) {
@@ -901,7 +952,6 @@ bool TTSTransformer::project_text_tokens(const int32_t * text_tokens, int32_t n_
ggml_free(ctx0);
return false;
}
struct ggml_tensor * out = ggml_graph_get_tensor(gf, "text_proj_out");
if (!out) {
error_msg_ = "Failed to find text projection output tensor";
@@ -922,7 +972,10 @@ bool TTSTransformer::build_prefill_graph(const int32_t * text_tokens, int32_t n_
const float * speaker_embd, int32_t language_id,
std::vector<float> & prefill_embd,
std::vector<float> & trailing_text_hidden,
std::vector<float> & tts_pad_embed) {
std::vector<float> & tts_pad_embed,
int32_t speaker_token_id,
const int32_t * instruct_tokens,
int32_t n_instruct_tokens) {
if (!text_tokens) {
error_msg_ = "text_tokens is null";
return false;
@@ -990,7 +1043,18 @@ bool TTSTransformer::build_prefill_graph(const int32_t * text_tokens, int32_t n_
return false;
}
const bool has_speaker = (speaker_embd != nullptr);
// If speaker_token_id is set, look up codec_embd[speaker_token_id] as the speaker embedding
std::vector<float> speaker_token_embed;
const float * effective_speaker_embd = speaker_embd;
if (speaker_token_id >= 0) {
speaker_token_embed.resize(hidden_size);
if (!lookup_single_embedding_row(model_.codec_embd, speaker_token_id, speaker_token_embed.data())) {
return false;
}
effective_speaker_embd = speaker_token_embed.data();
}
const bool has_speaker = (effective_speaker_embd != nullptr);
const int32_t codec_input_len = (int32_t)codec_prefill_tokens.size() + (has_speaker ? 1 : 0) + 2;
std::vector<float> codec_input_embedding((size_t)codec_input_len * hidden_size);
@@ -1000,7 +1064,7 @@ bool TTSTransformer::build_prefill_graph(const int32_t * text_tokens, int32_t n_
if (has_speaker) {
memcpy(codec_input_embedding.data() + (size_t)dst_token * hidden_size,
speaker_embd, hidden_size * sizeof(float));
effective_speaker_embd, hidden_size * sizeof(float));
++dst_token;
}
@@ -1031,12 +1095,33 @@ bool TTSTransformer::build_prefill_graph(const int32_t * text_tokens, int32_t n_
first_text_plus_codec_bos[h] = first_text_embed[h] + codec_bos_embed[h];
}
const int32_t prefill_len = 3 + codec_plus_overlay_len + 1;
// Project instruct tokens if provided
std::vector<float> instruct_proj;
if (instruct_tokens && n_instruct_tokens > 0) {
if (!project_text_tokens(instruct_tokens, n_instruct_tokens, instruct_proj)) {
return false;
}
}
const int32_t instruct_len = (int32_t)(instruct_proj.size() / hidden_size);
const int32_t prefill_len = instruct_len + 3 + codec_plus_overlay_len + 1;
prefill_embd.resize((size_t)prefill_len * hidden_size);
memcpy(prefill_embd.data(), role_embed.data(), role_embed.size() * sizeof(float));
memcpy(prefill_embd.data() + (size_t)3 * hidden_size,
codec_plus_overlay.data(), codec_plus_overlay.size() * sizeof(float));
memcpy(prefill_embd.data() + (size_t)(prefill_len - 1) * hidden_size,
int32_t pos = 0;
// instruct projection first (matches Python ordering: instruct before role)
if (instruct_len > 0) {
memcpy(prefill_embd.data(), instruct_proj.data(), (size_t)instruct_len * hidden_size * sizeof(float));
pos += instruct_len;
}
// role_embed (3 tokens)
memcpy(prefill_embd.data() + (size_t)pos * hidden_size,
role_embed.data(), (size_t)3 * hidden_size * sizeof(float));
pos += 3;
// codec_plus_overlay
memcpy(prefill_embd.data() + (size_t)pos * hidden_size,
codec_plus_overlay.data(), (size_t)codec_plus_overlay_len * hidden_size * sizeof(float));
pos += codec_plus_overlay_len;
// first_text + codec_bos (last token)
memcpy(prefill_embd.data() + (size_t)pos * hidden_size,
first_text_plus_codec_bos.data(), hidden_size * sizeof(float));
const int32_t trailing_token_count = std::max(0, n_tokens - 9);
@@ -1097,9 +1182,9 @@ struct ggml_cgraph * TTSTransformer::build_prefill_forward_graph(int32_t n_token
cur = ggml_rms_norm(ctx0, inpL, eps);
cur = ggml_mul(ctx0, cur, layer.attn_norm);
struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.attn_q, cur);
struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.attn_k, cur);
struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.attn_v, cur);
struct ggml_tensor * Qcur = mul_mat(ctx0, layer.attn_q, cur);
struct ggml_tensor * Kcur = mul_mat(ctx0, layer.attn_k, cur);
struct ggml_tensor * Vcur = mul_mat(ctx0, layer.attn_v, cur);
Qcur = ggml_reshape_3d(ctx0, Qcur, head_dim, n_head, n_tokens);
Kcur = ggml_reshape_3d(ctx0, Kcur, head_dim, n_kv_head, n_tokens);
@@ -1153,33 +1238,33 @@ struct ggml_cgraph * TTSTransformer::build_prefill_forward_graph(int32_t n_token
K = ggml_permute(ctx0, K, 0, 2, 1, 3);
V = ggml_permute(ctx0, V, 0, 2, 1, 3);
struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q);
struct ggml_tensor * KQ = mul_mat(ctx0, K, Q);
KQ = ggml_scale(ctx0, KQ, KQscale);
KQ = ggml_diag_mask_inf(ctx0, KQ, n_past);
KQ = ggml_soft_max(ctx0, KQ);
V = ggml_cont(ctx0, ggml_transpose(ctx0, V));
struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ);
struct ggml_tensor * KQV = mul_mat(ctx0, V, KQ);
KQV = ggml_permute(ctx0, KQV, 0, 2, 1, 3);
cur = ggml_cont_2d(ctx0, KQV, n_head * head_dim, n_tokens);
cur = ggml_mul_mat(ctx0, layer.attn_output, cur);
cur = mul_mat(ctx0, layer.attn_output, cur);
cur = ggml_add(ctx0, cur, inpL);
struct ggml_tensor * inpFF = cur;
cur = ggml_rms_norm(ctx0, inpFF, eps);
cur = ggml_mul(ctx0, cur, layer.ffn_norm);
struct ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ffn_gate, cur);
struct ggml_tensor * up = ggml_mul_mat(ctx0, layer.ffn_up, cur);
struct ggml_tensor * gate = mul_mat(ctx0, layer.ffn_gate, cur);
struct ggml_tensor * up = mul_mat(ctx0, layer.ffn_up, cur);
gate = ggml_silu(ctx0, gate);
cur = ggml_mul(ctx0, gate, up);
struct ggml_tensor * ffn_down_f32 = ggml_cast(ctx0, layer.ffn_down, GGML_TYPE_F32);
cur = ggml_mul_mat(ctx0, ffn_down_f32, cur);
cur = mul_mat(ctx0, ffn_down_f32, cur);
inpL = ggml_add(ctx0, cur, inpFF);
}
@@ -1191,7 +1276,7 @@ struct ggml_cgraph * TTSTransformer::build_prefill_forward_graph(int32_t n_token
ggml_set_name(cur, "hidden_states");
ggml_set_output(cur);
struct ggml_tensor * logits = ggml_mul_mat(ctx0, model_.codec_head, cur);
struct ggml_tensor * logits = mul_mat(ctx0, model_.codec_head, cur);
ggml_set_name(logits, "logits");
ggml_set_output(logits);
@@ -1242,9 +1327,9 @@ struct ggml_cgraph * TTSTransformer::build_step_graph(int32_t n_past) {
cur = ggml_rms_norm(ctx0, inpL, eps);
cur = ggml_mul(ctx0, cur, layer.attn_norm);
struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.attn_q, cur);
struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.attn_k, cur);
struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.attn_v, cur);
struct ggml_tensor * Qcur = mul_mat(ctx0, layer.attn_q, cur);
struct ggml_tensor * Kcur = mul_mat(ctx0, layer.attn_k, cur);
struct ggml_tensor * Vcur = mul_mat(ctx0, layer.attn_v, cur);
Qcur = ggml_reshape_3d(ctx0, Qcur, head_dim, n_head, n_tokens);
Kcur = ggml_reshape_3d(ctx0, Kcur, head_dim, n_kv_head, n_tokens);
@@ -1298,33 +1383,33 @@ struct ggml_cgraph * TTSTransformer::build_step_graph(int32_t n_past) {
K = ggml_permute(ctx0, K, 0, 2, 1, 3);
V = ggml_permute(ctx0, V, 0, 2, 1, 3);
struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q);
struct ggml_tensor * KQ = mul_mat(ctx0, K, Q);
KQ = ggml_scale(ctx0, KQ, KQscale);
KQ = ggml_diag_mask_inf(ctx0, KQ, n_past);
KQ = ggml_soft_max(ctx0, KQ);
V = ggml_cont(ctx0, ggml_transpose(ctx0, V));
struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ);
struct ggml_tensor * KQV = mul_mat(ctx0, V, KQ);
KQV = ggml_permute(ctx0, KQV, 0, 2, 1, 3);
cur = ggml_cont_2d(ctx0, KQV, n_head * head_dim, n_tokens);
cur = ggml_mul_mat(ctx0, layer.attn_output, cur);
cur = mul_mat(ctx0, layer.attn_output, cur);
cur = ggml_add(ctx0, cur, inpL);
struct ggml_tensor * inpFF = cur;
cur = ggml_rms_norm(ctx0, inpFF, eps);
cur = ggml_mul(ctx0, cur, layer.ffn_norm);
struct ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ffn_gate, cur);
struct ggml_tensor * up = ggml_mul_mat(ctx0, layer.ffn_up, cur);
struct ggml_tensor * gate = mul_mat(ctx0, layer.ffn_gate, cur);
struct ggml_tensor * up = mul_mat(ctx0, layer.ffn_up, cur);
gate = ggml_silu(ctx0, gate);
cur = ggml_mul(ctx0, gate, up);
struct ggml_tensor * ffn_down_f32 = ggml_cast(ctx0, layer.ffn_down, GGML_TYPE_F32);
cur = ggml_mul_mat(ctx0, ffn_down_f32, cur);
cur = mul_mat(ctx0, ffn_down_f32, cur);
inpL = ggml_add(ctx0, cur, inpFF);
}
@@ -1336,7 +1421,7 @@ struct ggml_cgraph * TTSTransformer::build_step_graph(int32_t n_past) {
ggml_set_name(cur, "hidden_states");
ggml_set_output(cur);
struct ggml_tensor * logits = ggml_mul_mat(ctx0, model_.codec_head, cur);
struct ggml_tensor * logits = mul_mat(ctx0, model_.codec_head, cur);
ggml_set_name(logits, "logits");
ggml_set_output(logits);
@@ -1349,10 +1434,11 @@ struct ggml_cgraph * TTSTransformer::build_step_graph(int32_t n_past) {
struct ggml_cgraph * TTSTransformer::build_code_pred_graph(int32_t n_prev_codes) {
const auto & cfg = model_.config;
const int n_head = cfg.n_attention_heads;
const int n_kv_head = cfg.n_key_value_heads;
const int head_dim = cfg.head_dim;
const int hidden_size = cfg.hidden_size;
const int n_head = cfg.code_pred_n_attention_heads;
const int n_kv_head = cfg.code_pred_n_key_value_heads;
const int head_dim = cfg.code_pred_head_dim;
const int cp_hidden = cfg.code_pred_hidden_size;
const int talker_hidden = cfg.hidden_size;
const float eps = cfg.rms_norm_eps;
const int n_layer = cfg.code_pred_layers;
const int n_codebooks = cfg.n_codebooks;
@@ -1366,7 +1452,7 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_graph(int32_t n_prev_codes)
struct ggml_context * ctx0 = ggml_init(params);
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false);
struct ggml_tensor * inp_hidden = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hidden_size);
struct ggml_tensor * inp_hidden = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, talker_hidden);
ggml_set_name(inp_hidden, "inp_hidden");
ggml_set_input(inp_hidden);
@@ -1377,12 +1463,27 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_graph(int32_t n_prev_codes)
ggml_set_input(inp_prev_codes);
}
struct ggml_tensor * cur = ggml_reshape_2d(ctx0, inp_hidden, hidden_size, 1);
// Project talker hidden to code_pred hidden
struct ggml_tensor * cur = ggml_reshape_2d(ctx0, inp_hidden, talker_hidden, 1);
if (model_.code_pred_mtp_proj_w) {
cur = mul_mat(ctx0, model_.code_pred_mtp_proj_w, cur);
if (model_.code_pred_mtp_proj_b) {
cur = ggml_add(ctx0, cur, model_.code_pred_mtp_proj_b);
}
}
if (n_prev_codes > 0 && inp_prev_codes) {
for (int cb = 0; cb < n_prev_codes && cb < n_codebooks - 1; ++cb) {
struct ggml_tensor * code_idx = ggml_view_1d(ctx0, inp_prev_codes, 1, cb * sizeof(int32_t));
struct ggml_tensor * code_embd = ggml_get_rows(ctx0, model_.code_pred_embd[cb], code_idx);
struct ggml_tensor * code_embd_raw = ggml_get_rows(ctx0, model_.code_pred_embd[cb], code_idx);
// code_pred_embd is [talker_hidden, vocab] — project down
struct ggml_tensor * code_embd = code_embd_raw;
if (model_.code_pred_mtp_proj_w) {
code_embd = mul_mat(ctx0, model_.code_pred_mtp_proj_w, code_embd_raw);
if (model_.code_pred_mtp_proj_b) {
code_embd = ggml_add(ctx0, code_embd, model_.code_pred_mtp_proj_b);
}
}
cur = ggml_add(ctx0, cur, code_embd);
}
}
@@ -1397,9 +1498,9 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_graph(int32_t n_prev_codes)
cur = ggml_rms_norm(ctx0, inpL, eps);
cur = ggml_mul(ctx0, cur, layer.attn_norm);
struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.attn_q, cur);
struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.attn_k, cur);
struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.attn_v, cur);
struct ggml_tensor * Qcur = mul_mat(ctx0, layer.attn_q, cur);
struct ggml_tensor * Kcur = mul_mat(ctx0, layer.attn_k, cur);
struct ggml_tensor * Vcur = mul_mat(ctx0, layer.attn_v, cur);
Qcur = ggml_reshape_3d(ctx0, Qcur, head_dim, n_head, 1);
Kcur = ggml_reshape_3d(ctx0, Kcur, head_dim, n_kv_head, 1);
@@ -1419,41 +1520,45 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_graph(int32_t n_prev_codes)
struct ggml_tensor * K = ggml_permute(ctx0, Kcur, 0, 2, 1, 3);
struct ggml_tensor * V = ggml_permute(ctx0, Vcur, 0, 2, 1, 3);
struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q);
struct ggml_tensor * KQ = mul_mat(ctx0, K, Q);
KQ = ggml_scale(ctx0, KQ, KQscale);
KQ = ggml_soft_max(ctx0, KQ);
V = ggml_cont(ctx0, ggml_transpose(ctx0, V));
struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ);
struct ggml_tensor * KQV = mul_mat(ctx0, V, KQ);
KQV = ggml_permute(ctx0, KQV, 0, 2, 1, 3);
cur = ggml_cont_2d(ctx0, KQV, n_head * head_dim, 1);
cur = ggml_mul_mat(ctx0, layer.attn_output, cur);
cur = mul_mat(ctx0, layer.attn_output, cur);
cur = ggml_add(ctx0, cur, inpL);
struct ggml_tensor * inpFF = cur;
cur = ggml_rms_norm(ctx0, inpFF, eps);
cur = ggml_mul(ctx0, cur, layer.ffn_norm);
struct ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ffn_gate, cur);
struct ggml_tensor * up = ggml_mul_mat(ctx0, layer.ffn_up, cur);
struct ggml_tensor * gate = mul_mat(ctx0, layer.ffn_gate, cur);
struct ggml_tensor * up = mul_mat(ctx0, layer.ffn_up, cur);
gate = ggml_silu(ctx0, gate);
cur = ggml_mul(ctx0, gate, up);
struct ggml_tensor * old_ffn_down_f32 = ggml_cast(ctx0, layer.ffn_down, GGML_TYPE_F32);
cur = ggml_mul_mat(ctx0, old_ffn_down_f32, cur);
cur = mul_mat(ctx0, old_ffn_down_f32, cur);
inpL = ggml_add(ctx0, cur, inpFF);
}
cur = inpL;
// Apply output normalization (matching prefill/step paths)
cur = ggml_rms_norm(ctx0, cur, eps);
cur = ggml_mul(ctx0, cur, model_.code_pred_output_norm);
std::vector<struct ggml_tensor *> all_logits;
for (int cb = 0; cb < n_codebooks - 1; ++cb) {
struct ggml_tensor * cb_logits = ggml_mul_mat(ctx0, model_.code_pred_head[cb], cur);
struct ggml_tensor * cb_logits = mul_mat(ctx0, model_.code_pred_head[cb], cur);
ggml_format_name(cb_logits, "logits_cb%d", cb + 1);
ggml_set_output(cb_logits);
all_logits.push_back(cb_logits);
@@ -1470,10 +1575,11 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_graph(int32_t n_prev_codes)
struct ggml_cgraph * TTSTransformer::build_code_pred_prefill_graph() {
const auto & cfg = model_.config;
const int n_head = cfg.n_attention_heads;
const int n_kv_head = cfg.n_key_value_heads;
const int head_dim = cfg.head_dim;
const int hidden_size = cfg.hidden_size;
const int n_head = cfg.code_pred_n_attention_heads;
const int n_kv_head = cfg.code_pred_n_key_value_heads;
const int head_dim = cfg.code_pred_head_dim;
const int cp_hidden = cfg.code_pred_hidden_size;
const int talker_hidden = cfg.hidden_size;
const float eps = cfg.rms_norm_eps;
const float rope_theta = cfg.rope_theta;
const int n_layer = cfg.code_pred_layers;
@@ -1488,13 +1594,13 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_prefill_graph() {
struct ggml_context * ctx0 = ggml_init(params);
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false);
// Input: past_hidden from talker [hidden_size]
struct ggml_tensor * inp_hidden = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hidden_size);
// Input: past_hidden from talker [talker_hidden]
struct ggml_tensor * inp_hidden = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, talker_hidden);
ggml_set_name(inp_hidden, "inp_hidden");
ggml_set_input(inp_hidden);
// Input: codebook 0 token embedding [hidden_size] (pre-computed using talker's codec_embd)
struct ggml_tensor * inp_cb0_embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hidden_size);
// Input: codebook 0 token embedding [talker_hidden] (pre-computed using talker's codec_embd)
struct ggml_tensor * inp_cb0_embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, talker_hidden);
ggml_set_name(inp_cb0_embd, "inp_cb0_embd");
ggml_set_input(inp_cb0_embd);
@@ -1502,11 +1608,19 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_prefill_graph() {
ggml_set_name(inp_pos, "inp_pos");
ggml_set_input(inp_pos);
// Concatenate [past_hidden, cb0_embd] -> [2, hidden_size]
struct ggml_tensor * hidden_2d = ggml_reshape_2d(ctx0, inp_hidden, hidden_size, 1);
struct ggml_tensor * cb0_2d = ggml_reshape_2d(ctx0, inp_cb0_embd, hidden_size, 1);
// Concatenate [past_hidden, cb0_embd] -> [talker_hidden, 2]
struct ggml_tensor * hidden_2d = ggml_reshape_2d(ctx0, inp_hidden, talker_hidden, 1);
struct ggml_tensor * cb0_2d = ggml_reshape_2d(ctx0, inp_cb0_embd, talker_hidden, 1);
struct ggml_tensor * cur = ggml_concat(ctx0, hidden_2d, cb0_2d, 1);
// Apply MTP projection if present (talker_hidden -> code_pred_hidden)
if (model_.code_pred_mtp_proj_w) {
cur = mul_mat(ctx0, model_.code_pred_mtp_proj_w, cur);
if (model_.code_pred_mtp_proj_b) {
cur = ggml_add(ctx0, cur, model_.code_pred_mtp_proj_b);
}
}
struct ggml_tensor * inpL = cur;
const float KQscale = 1.0f / sqrtf(float(head_dim));
@@ -1517,9 +1631,9 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_prefill_graph() {
cur = ggml_rms_norm(ctx0, inpL, eps);
cur = ggml_mul(ctx0, cur, layer.attn_norm);
struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.attn_q, cur);
struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.attn_k, cur);
struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.attn_v, cur);
struct ggml_tensor * Qcur = mul_mat(ctx0, layer.attn_q, cur);
struct ggml_tensor * Kcur = mul_mat(ctx0, layer.attn_k, cur);
struct ggml_tensor * Vcur = mul_mat(ctx0, layer.attn_v, cur);
Qcur = ggml_reshape_3d(ctx0, Qcur, head_dim, n_head, n_tokens);
Kcur = ggml_reshape_3d(ctx0, Kcur, head_dim, n_kv_head, n_tokens);
@@ -1562,33 +1676,33 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_prefill_graph() {
struct ggml_tensor * K = ggml_permute(ctx0, Kcur, 0, 2, 1, 3);
struct ggml_tensor * V = ggml_permute(ctx0, Vcur, 0, 2, 1, 3);
struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q);
struct ggml_tensor * KQ = mul_mat(ctx0, K, Q);
KQ = ggml_scale(ctx0, KQ, KQscale);
KQ = ggml_diag_mask_inf(ctx0, KQ, 0);
KQ = ggml_soft_max(ctx0, KQ);
V = ggml_cont(ctx0, ggml_transpose(ctx0, V));
struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ);
struct ggml_tensor * KQV = mul_mat(ctx0, V, KQ);
KQV = ggml_permute(ctx0, KQV, 0, 2, 1, 3);
cur = ggml_cont_2d(ctx0, KQV, n_head * head_dim, n_tokens);
cur = ggml_mul_mat(ctx0, layer.attn_output, cur);
cur = mul_mat(ctx0, layer.attn_output, cur);
cur = ggml_add(ctx0, cur, inpL);
struct ggml_tensor * inpFF = cur;
cur = ggml_rms_norm(ctx0, inpFF, eps);
cur = ggml_mul(ctx0, cur, layer.ffn_norm);
struct ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ffn_gate, cur);
struct ggml_tensor * up = ggml_mul_mat(ctx0, layer.ffn_up, cur);
struct ggml_tensor * gate = mul_mat(ctx0, layer.ffn_gate, cur);
struct ggml_tensor * up = mul_mat(ctx0, layer.ffn_up, cur);
gate = ggml_silu(ctx0, gate);
cur = ggml_mul(ctx0, gate, up);
struct ggml_tensor * ffn_down_f32 = ggml_cast(ctx0, layer.ffn_down, GGML_TYPE_F32);
cur = ggml_mul_mat(ctx0, ffn_down_f32, cur);
cur = mul_mat(ctx0, ffn_down_f32, cur);
inpL = ggml_add(ctx0, cur, inpFF);
}
@@ -1598,10 +1712,10 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_prefill_graph() {
cur = ggml_rms_norm(ctx0, cur, eps);
cur = ggml_mul(ctx0, cur, model_.code_pred_output_norm);
struct ggml_tensor * last_hidden = ggml_view_2d(ctx0, cur, hidden_size, 1,
cur->nb[1], hidden_size * sizeof(float));
struct ggml_tensor * last_hidden = ggml_view_2d(ctx0, cur, cp_hidden, 1,
cur->nb[1], cp_hidden * sizeof(float));
struct ggml_tensor * logits = ggml_mul_mat(ctx0, model_.code_pred_head[0], last_hidden);
struct ggml_tensor * logits = mul_mat(ctx0, model_.code_pred_head[0], last_hidden);
ggml_set_name(logits, "logits");
ggml_set_output(logits);
@@ -1614,10 +1728,11 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_prefill_graph() {
struct ggml_cgraph * TTSTransformer::build_code_pred_step_graph(int32_t n_past, int32_t generation_step) {
const auto & cfg = model_.config;
const int n_head = cfg.n_attention_heads;
const int n_kv_head = cfg.n_key_value_heads;
const int head_dim = cfg.head_dim;
const int hidden_size = cfg.hidden_size;
const int n_head = cfg.code_pred_n_attention_heads;
const int n_kv_head = cfg.code_pred_n_key_value_heads;
const int head_dim = cfg.code_pred_head_dim;
const int cp_hidden = cfg.code_pred_hidden_size;
const int talker_hidden = cfg.hidden_size;
const float eps = cfg.rms_norm_eps;
const float rope_theta = cfg.rope_theta;
const int n_layer = cfg.code_pred_layers;
@@ -1632,7 +1747,8 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_step_graph(int32_t n_past,
struct ggml_context * ctx0 = ggml_init(params);
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, QWEN3_TTS_MAX_NODES, false);
struct ggml_tensor * inp_hidden = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hidden_size);
// inp_hidden is only used when generation_step == 0 (not used in step graph normally)
struct ggml_tensor * inp_hidden = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, talker_hidden);
ggml_set_name(inp_hidden, "inp_hidden");
ggml_set_input(inp_hidden);
@@ -1646,10 +1762,24 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_step_graph(int32_t n_past,
struct ggml_tensor * cur;
if (generation_step == 0) {
cur = ggml_reshape_2d(ctx0, inp_hidden, hidden_size, 1);
// inp_hidden is talker-dim, project down
cur = ggml_reshape_2d(ctx0, inp_hidden, talker_hidden, 1);
if (model_.code_pred_mtp_proj_w) {
cur = mul_mat(ctx0, model_.code_pred_mtp_proj_w, cur);
if (model_.code_pred_mtp_proj_b) {
cur = ggml_add(ctx0, cur, model_.code_pred_mtp_proj_b);
}
}
} else {
// code_pred_embd is [talker_hidden, vocab] — lookup returns talker-dim, project down
cur = ggml_get_rows(ctx0, model_.code_pred_embd[generation_step - 1], inp_code);
cur = ggml_reshape_2d(ctx0, cur, hidden_size, 1);
cur = ggml_reshape_2d(ctx0, cur, talker_hidden, 1);
if (model_.code_pred_mtp_proj_w) {
cur = mul_mat(ctx0, model_.code_pred_mtp_proj_w, cur);
if (model_.code_pred_mtp_proj_b) {
cur = ggml_add(ctx0, cur, model_.code_pred_mtp_proj_b);
}
}
}
struct ggml_tensor * inpL = cur;
@@ -1662,9 +1792,9 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_step_graph(int32_t n_past,
cur = ggml_rms_norm(ctx0, inpL, eps);
cur = ggml_mul(ctx0, cur, layer.attn_norm);
struct ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.attn_q, cur);
struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.attn_k, cur);
struct ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.attn_v, cur);
struct ggml_tensor * Qcur = mul_mat(ctx0, layer.attn_q, cur);
struct ggml_tensor * Kcur = mul_mat(ctx0, layer.attn_k, cur);
struct ggml_tensor * Vcur = mul_mat(ctx0, layer.attn_v, cur);
Qcur = ggml_reshape_3d(ctx0, Qcur, head_dim, n_head, n_tokens);
Kcur = ggml_reshape_3d(ctx0, Kcur, head_dim, n_kv_head, n_tokens);
@@ -1718,33 +1848,33 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_step_graph(int32_t n_past,
K = ggml_permute(ctx0, K, 0, 2, 1, 3);
V = ggml_permute(ctx0, V, 0, 2, 1, 3);
struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q);
struct ggml_tensor * KQ = mul_mat(ctx0, K, Q);
KQ = ggml_scale(ctx0, KQ, KQscale);
KQ = ggml_diag_mask_inf(ctx0, KQ, n_past);
KQ = ggml_soft_max(ctx0, KQ);
V = ggml_cont(ctx0, ggml_transpose(ctx0, V));
struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ);
struct ggml_tensor * KQV = mul_mat(ctx0, V, KQ);
KQV = ggml_permute(ctx0, KQV, 0, 2, 1, 3);
cur = ggml_cont_2d(ctx0, KQV, n_head * head_dim, n_tokens);
cur = ggml_mul_mat(ctx0, layer.attn_output, cur);
cur = mul_mat(ctx0, layer.attn_output, cur);
cur = ggml_add(ctx0, cur, inpL);
struct ggml_tensor * inpFF = cur;
cur = ggml_rms_norm(ctx0, inpFF, eps);
cur = ggml_mul(ctx0, cur, layer.ffn_norm);
struct ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ffn_gate, cur);
struct ggml_tensor * up = ggml_mul_mat(ctx0, layer.ffn_up, cur);
struct ggml_tensor * gate = mul_mat(ctx0, layer.ffn_gate, cur);
struct ggml_tensor * up = mul_mat(ctx0, layer.ffn_up, cur);
gate = ggml_silu(ctx0, gate);
cur = ggml_mul(ctx0, gate, up);
struct ggml_tensor * step_ffn_down_f32 = ggml_cast(ctx0, layer.ffn_down, GGML_TYPE_F32);
cur = ggml_mul_mat(ctx0, step_ffn_down_f32, cur);
cur = mul_mat(ctx0, step_ffn_down_f32, cur);
inpL = ggml_add(ctx0, cur, inpFF);
}
@@ -1754,7 +1884,7 @@ struct ggml_cgraph * TTSTransformer::build_code_pred_step_graph(int32_t n_past,
cur = ggml_rms_norm(ctx0, cur, eps);
cur = ggml_mul(ctx0, cur, model_.code_pred_output_norm);
struct ggml_tensor * logits = ggml_mul_mat(ctx0, model_.code_pred_head[generation_step], cur);
struct ggml_tensor * logits = mul_mat(ctx0, model_.code_pred_head[generation_step], cur);
ggml_set_name(logits, "logits");
ggml_set_output(logits);
@@ -2010,17 +2140,20 @@ bool TTSTransformer::forward_step(const float * step_embd, int32_t n_past,
#endif
struct ggml_tensor * hidden = ggml_graph_get_tensor(gf, "hidden_states");
if (!hidden) {
error_msg_ = "Failed to find hidden_states tensor in step graph";
ggml_backend_sched_reset(state_.sched);
return false;
}
#ifdef QWEN3_TTS_TIMING
t0 = clk::now();
#endif
if (hidden) {
last_hidden_.resize(model_.config.hidden_size);
ggml_backend_tensor_get(hidden, last_hidden_.data(), 0,
model_.config.hidden_size * sizeof(float));
if (hidden_out) {
*hidden_out = last_hidden_;
}
last_hidden_.resize(model_.config.hidden_size);
ggml_backend_tensor_get(hidden, last_hidden_.data(), 0,
model_.config.hidden_size * sizeof(float));
if (hidden_out) {
*hidden_out = last_hidden_;
}
struct ggml_tensor * logits = ggml_graph_get_tensor(gf, "logits");
@@ -2144,7 +2277,7 @@ bool TTSTransformer::predict_codes_autoregressive_coreml(const float * hidden,
output.resize(n_steps);
std::vector<float> logits_data(cfg.code_pred_vocab_size);
std::vector<float> code_probs(cfg.code_pred_vocab_size);
std::vector<float> seq_embd((size_t)16 * cfg.hidden_size, 0.0f);
std::vector<float> seq_embd((size_t)cfg.n_codebooks * cfg.hidden_size, 0.0f);
#ifdef QWEN3_TTS_TIMING
using clk = std::chrono::high_resolution_clock;
@@ -2270,14 +2403,23 @@ bool TTSTransformer::predict_codes_autoregressive(const float * hidden, int32_t
use_coreml_code_predictor_ = false;
}
if (state_.code_pred_cache.n_ctx < 16) {
if (!init_code_pred_kv_cache(16)) {
if (!model_.code_pred_output_norm) {
error_msg_ = "code_pred_output_norm not loaded (required for GGML code predictor)";
return false;
}
if (model_.code_pred_head.empty()) {
error_msg_ = "code_pred_head tensors not loaded (required for GGML code predictor)";
return false;
}
if (state_.code_pred_cache.n_ctx < cfg.n_codebooks) {
if (!init_code_pred_kv_cache(cfg.n_codebooks)) {
return false;
}
}
clear_code_pred_kv_cache();
output.resize(15);
output.resize(cfg.n_codebooks - 1);
std::vector<float> logits_data(cfg.code_pred_vocab_size);
std::vector<float> code_probs(cfg.code_pred_vocab_size);
@@ -2422,7 +2564,7 @@ bool TTSTransformer::predict_codes_autoregressive(const float * hidden, int32_t
#ifdef QWEN3_TTS_TIMING
auto t_steps_start = clk::now();
#endif
for (int step = 1; step < 15; ++step) {
for (int step = 1; step < cfg.n_codebooks - 1; ++step) {
int32_t n_past = step + 1;
#ifdef QWEN3_TTS_TIMING
@@ -2517,7 +2659,11 @@ bool TTSTransformer::generate(const int32_t * text_tokens, int32_t n_tokens,
int32_t language_id,
float repetition_penalty,
float temperature,
int32_t top_k) {
int32_t top_k,
int32_t speaker_token_id,
const int32_t * instruct_tokens,
int32_t n_instruct_tokens,
std::function<void(int, int)> progress_cb) {
#ifdef QWEN3_TTS_TIMING
using clk = std::chrono::high_resolution_clock;
tts_timing timing = {};
@@ -2553,7 +2699,8 @@ bool TTSTransformer::generate(const int32_t * text_tokens, int32_t n_tokens,
t0 = clk::now();
#endif
if (!build_prefill_graph(text_tokens, n_tokens, speaker_embd, language_id,
prefill_embd, trailing_text_hidden, tts_pad_embed)) {
prefill_embd, trailing_text_hidden, tts_pad_embed,
speaker_token_id, instruct_tokens, n_instruct_tokens)) {
return false;
}
#ifdef QWEN3_TTS_TIMING
@@ -2571,7 +2718,6 @@ bool TTSTransformer::generate(const int32_t * text_tokens, int32_t n_tokens,
}
}
clear_kv_cache();
std::vector<float> hidden_out;
std::vector<float> logits;
@@ -2592,7 +2738,7 @@ bool TTSTransformer::generate(const int32_t * text_tokens, int32_t n_tokens,
int32_t n_past = prefill_len;
std::vector<int32_t> frame_codes(cfg.n_codebooks);
std::unordered_set<int32_t> generated_cb0_tokens;
const int32_t suppress_start = cfg.codec_vocab_size - 1024;
const int32_t suppress_start = (cfg.codec_vocab_size >= 1024) ? (cfg.codec_vocab_size - 1024) : 0;
std::vector<float> probs(cfg.codec_vocab_size);
std::vector<float> step_embd(cfg.hidden_size, 0.0f);
@@ -2685,6 +2831,10 @@ bool TTSTransformer::generate(const int32_t * text_tokens, int32_t n_tokens,
output.push_back(frame_codes[cb]);
}
if (progress_cb) {
progress_cb(frame + 1, max_len);
}
#ifdef QWEN3_TTS_TIMING
timing.n_frames = frame + 1;
#endif
otherarch/qwen3tts/tts_transformer.h
+35 -4
View File
@@ -10,6 +10,7 @@
#include <vector>
#include <memory>
#include <random>
#include <functional>
#ifdef QWEN3_TTS_TIMING
#include <chrono>
#endif
@@ -80,6 +81,11 @@ struct tts_transformer_config {
// Code predictor
int32_t code_pred_layers = 5;
int32_t code_pred_vocab_size = 2048; // Per-codebook vocab
int32_t code_pred_hidden_size = 1024;
int32_t code_pred_n_attention_heads = 16;
int32_t code_pred_n_key_value_heads = 8;
int32_t code_pred_intermediate_size = 3072;
int32_t code_pred_head_dim = 128;
// Special codec tokens
int32_t codec_pad_id = 2148;
@@ -143,7 +149,11 @@ struct tts_transformer_model {
std::vector<transformer_layer> code_pred_layers;
// Code predictor output norm (final RMS norm before lm_head)
struct ggml_tensor * code_pred_output_norm = nullptr; // [hidden_size]
struct ggml_tensor * code_pred_output_norm = nullptr; // [code_pred_hidden_size]
// Code predictor MTP projection (talker_hidden -> code_pred_hidden)
struct ggml_tensor * code_pred_mtp_proj_w = nullptr; // [hidden_size, code_pred_hidden_size]
struct ggml_tensor * code_pred_mtp_proj_b = nullptr; // [code_pred_hidden_size]
// Code predictor per-codebook embeddings and heads (15 codebooks, 0 uses talker output)
std::vector<struct ggml_tensor *> code_pred_embd; // [hidden_size, code_pred_vocab_size] x 15
@@ -260,18 +270,29 @@ public:
// speaker_embd: speaker embedding [hidden_size]
// max_len: maximum number of frames to generate
// output: generated speech codes [n_frames, n_codebooks]
// Progress callback: (current_frame, max_frames) called after each frame
using progress_callback_t = std::function<void(int, int)>;
bool generate(const int32_t * text_tokens, int32_t n_tokens,
const float * speaker_embd, int32_t max_len,
std::vector<int32_t> & output,
int32_t language_id = 2050,
float repetition_penalty = 1.05f,
float temperature = 0.9f,
int32_t top_k = 50);
int32_t top_k = 50,
int32_t speaker_token_id = -1,
const int32_t * instruct_tokens = nullptr,
int32_t n_instruct_tokens = 0,
progress_callback_t progress_cb = nullptr);
const tts_transformer_config & get_config() const { return model_.config; }
const std::string & get_error() const { return error_msg_; }
void set_force_f32_acc(bool v) { force_f32_acc_ = v; }
void set_force_cpu(bool v) { force_cpu_ = v; }
void set_seed(uint32_t seed) { if (seed) rng_.seed(seed); }
// Legacy interface for compatibility
bool forward(const int32_t * tokens, int32_t n_tokens, int32_t n_past,
std::vector<float> & output);
@@ -292,7 +313,10 @@ private:
const float * speaker_embd, int32_t language_id,
std::vector<float> & prefill_embd,
std::vector<float> & trailing_text_hidden,
std::vector<float> & tts_pad_embed);
std::vector<float> & tts_pad_embed,
int32_t speaker_token_id = -1,
const int32_t * instruct_tokens = nullptr,
int32_t n_instruct_tokens = 0);
struct ggml_cgraph * build_prefill_forward_graph(int32_t n_tokens, int32_t n_past);
@@ -328,9 +352,16 @@ private:
// Load tensor data from file
bool load_tensor_data(const std::string & path, struct gguf_context * ctx);
// Helper: ggml_mul_mat with optional f32 accumulation precision
struct ggml_tensor * mul_mat(struct ggml_context * ctx,
struct ggml_tensor * a,
struct ggml_tensor * b);
tts_transformer_model model_;
tts_transformer_state state_;
std::string error_msg_;
bool force_f32_acc_ = false;
bool force_cpu_ = false;
// Cached hidden states from last forward pass
std::vector<float> last_hidden_;
@@ -352,4 +383,4 @@ void free_transformer_model(tts_transformer_model & model);
// Free KV cache resources
void free_tts_kv_cache(tts_kv_cache & cache);
} // namespace qwen3_tts
} // namespace qwen3_tts