247 lines
10 KiB
C++
247 lines
10 KiB
C++
#pragma once
|
|
// encoder-transformer.h: 8-layer Mimi-style transformer for the Qwen3-TTS
|
|
// encoder. Operates on the SEANet output stream at 25 Hz, 512 channels.
|
|
//
|
|
// Differs from the Qwen3-style decoder transformer on several points:
|
|
// - LayerNorm with bias (not RMSNorm)
|
|
// - Plain MLP fc1 -> GELU -> fc2 (not SwiGLU)
|
|
// - No biases on q/k/v/o projections
|
|
// - 8 attention heads instead of 16
|
|
// - intermediate_size 2048 instead of 1024
|
|
// - Pure causal attention, full T x T mask (the upstream config carries
|
|
// a sliding_window field but Mimi never applies it)
|
|
// - No top-level input_proj / output_proj brackets: the SEANet output
|
|
// already has hidden_size channels
|
|
//
|
|
// Common with the decoder side:
|
|
// - RoPE NEOX style with theta 10000
|
|
// - LayerScale per channel post-attention and post-MLP
|
|
// - Pre-norm residual on both attention and MLP
|
|
|
|
#include "ggml-backend.h"
|
|
#include "ggml.h"
|
|
#include "gguf-weights.h"
|
|
#include "weight-ctx.h"
|
|
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#define ENC_TRANS_MAX_LAYERS 16
|
|
|
|
struct QwenEncoderTransformerLayer {
|
|
// Pre-attention LayerNorm
|
|
struct ggml_tensor * input_norm_w;
|
|
struct ggml_tensor * input_norm_b;
|
|
// Attention projections, no bias
|
|
struct ggml_tensor * q_proj_w;
|
|
struct ggml_tensor * k_proj_w;
|
|
struct ggml_tensor * v_proj_w;
|
|
struct ggml_tensor * o_proj_w;
|
|
// Post-attention LayerScale per channel
|
|
struct ggml_tensor * attn_scale;
|
|
// Post-attention pre-MLP LayerNorm
|
|
struct ggml_tensor * post_attn_norm_w;
|
|
struct ggml_tensor * post_attn_norm_b;
|
|
// MLP
|
|
struct ggml_tensor * fc1_w;
|
|
struct ggml_tensor * fc2_w;
|
|
// Post-MLP LayerScale per channel
|
|
struct ggml_tensor * mlp_scale;
|
|
};
|
|
|
|
struct QwenEncoderTransformer {
|
|
int hidden_size;
|
|
int num_layers;
|
|
int num_attention_heads;
|
|
int num_kv_heads;
|
|
int head_dim;
|
|
int intermediate_size;
|
|
float rope_theta;
|
|
float norm_eps;
|
|
|
|
QwenEncoderTransformerLayer layers[ENC_TRANS_MAX_LAYERS];
|
|
|
|
struct ggml_context * weight_ctx;
|
|
ggml_backend_buffer_t weight_buf;
|
|
};
|
|
|
|
static bool enc_trans_load(QwenEncoderTransformer * tr, const GGUFModel & gf, ggml_backend_t backend) {
|
|
tr->hidden_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.hidden_size");
|
|
tr->num_layers = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.num_hidden_layers");
|
|
tr->num_attention_heads = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.num_attention_heads");
|
|
tr->num_kv_heads = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.num_key_value_heads");
|
|
tr->head_dim = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.head_dim");
|
|
tr->intermediate_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.intermediate_size");
|
|
tr->rope_theta = gf_get_f32(gf, "qwen3-tts-tokenizer.encoder.rope_theta");
|
|
tr->norm_eps = gf_get_f32(gf, "qwen3-tts-tokenizer.encoder.norm_eps");
|
|
|
|
if (tr->num_layers > ENC_TRANS_MAX_LAYERS) {
|
|
fprintf(stderr, "[EncTransformer] FATAL: %d layers exceeds compile-time max %d\n", tr->num_layers,
|
|
ENC_TRANS_MAX_LAYERS);
|
|
return false;
|
|
}
|
|
|
|
int n_tensors = tr->num_layers * 12 + 4; // 12 tensors per layer + headroom
|
|
WeightCtx wctx;
|
|
wctx_init(&wctx, n_tensors);
|
|
|
|
for (int l = 0; l < tr->num_layers; l++) {
|
|
QwenEncoderTransformerLayer & ly = tr->layers[l];
|
|
char prefix[96];
|
|
snprintf(prefix, sizeof(prefix), "tok_enc.blk.%d", l);
|
|
std::string p(prefix);
|
|
|
|
ly.input_norm_w = gf_load_tensor(&wctx, gf, p + ".attn_norm.weight");
|
|
ly.input_norm_b = gf_load_tensor(&wctx, gf, p + ".attn_norm.bias");
|
|
ly.q_proj_w = gf_load_tensor(&wctx, gf, p + ".attn_q.weight");
|
|
ly.k_proj_w = gf_load_tensor(&wctx, gf, p + ".attn_k.weight");
|
|
ly.v_proj_w = gf_load_tensor(&wctx, gf, p + ".attn_v.weight");
|
|
ly.o_proj_w = gf_load_tensor(&wctx, gf, p + ".attn_output.weight");
|
|
ly.attn_scale = gf_load_tensor(&wctx, gf, p + ".attn_scale");
|
|
ly.post_attn_norm_w = gf_load_tensor(&wctx, gf, p + ".ffn_norm.weight");
|
|
ly.post_attn_norm_b = gf_load_tensor(&wctx, gf, p + ".ffn_norm.bias");
|
|
ly.fc1_w = gf_load_tensor(&wctx, gf, p + ".ffn_up.weight");
|
|
ly.fc2_w = gf_load_tensor(&wctx, gf, p + ".ffn_down.weight");
|
|
ly.mlp_scale = gf_load_tensor(&wctx, gf, p + ".ffn_scale");
|
|
}
|
|
|
|
if (!wctx_alloc(&wctx, backend)) {
|
|
fprintf(stderr, "[EncTransformer] FATAL: backend allocation failed\n");
|
|
return false;
|
|
}
|
|
tr->weight_ctx = wctx.ctx;
|
|
tr->weight_buf = wctx.buffer;
|
|
|
|
fprintf(stderr,
|
|
"[EncTransformer] Loaded: %d layers, hidden %d, heads %d/%d, head_dim %d, "
|
|
"FFN %d, RoPE theta %.0f\n",
|
|
tr->num_layers, tr->hidden_size, tr->num_attention_heads, tr->num_kv_heads, tr->head_dim,
|
|
tr->intermediate_size, tr->rope_theta);
|
|
return true;
|
|
}
|
|
|
|
static void enc_trans_free(QwenEncoderTransformer * tr) {
|
|
if (tr->weight_buf) {
|
|
ggml_backend_buffer_free(tr->weight_buf);
|
|
tr->weight_buf = NULL;
|
|
}
|
|
if (tr->weight_ctx) {
|
|
ggml_free(tr->weight_ctx);
|
|
tr->weight_ctx = NULL;
|
|
}
|
|
}
|
|
|
|
// Build a [T, T] additive causal mask (0 where allowed, -inf where masked).
|
|
// Pure causal: k <= q. The upstream config carries a sliding_window
|
|
// value but neither MimiAttention's eager forward nor MimiTransformerModel
|
|
// (create_causal_mask) ever apply it. The Qwen3TTS encoder inherits this
|
|
// convention, so we mirror it bit for bit here.
|
|
static void enc_trans_build_causal_mask(int T, std::vector<float> & dst) {
|
|
dst.assign((size_t) T * (size_t) T, -INFINITY);
|
|
for (int q = 0; q < T; q++) {
|
|
for (int k = 0; k <= q; k++) {
|
|
dst[(size_t) q * (size_t) T + (size_t) k] = 0.0f;
|
|
}
|
|
}
|
|
}
|
|
|
|
static void enc_trans_build_positions(int T, std::vector<int32_t> & dst) {
|
|
dst.resize((size_t) T);
|
|
for (int i = 0; i < T; i++) {
|
|
dst[i] = i;
|
|
}
|
|
}
|
|
|
|
// One Mimi transformer layer. Pre-LayerNorm with bias, attention without
|
|
// q/k/v/o biases, MLP with fc1 -> GELU -> fc2 (no SwiGLU), LayerScale on
|
|
// both residual paths.
|
|
// x : [hidden, T] f32 C-first
|
|
// positions: [T] i32
|
|
// mask : [T, T] f32 additive
|
|
// Returns [hidden, T] f32 C-first.
|
|
static struct ggml_tensor * enc_trans_layer_forward(struct ggml_context * ctx,
|
|
const QwenEncoderTransformer * tr,
|
|
const QwenEncoderTransformerLayer & layer,
|
|
struct ggml_tensor * x,
|
|
struct ggml_tensor * positions,
|
|
struct ggml_tensor * mask,
|
|
int T) {
|
|
int hidden = tr->hidden_size;
|
|
int n_q_heads = tr->num_attention_heads;
|
|
int n_kv = tr->num_kv_heads;
|
|
int hd = tr->head_dim;
|
|
|
|
// Pre-LayerNorm with affine (weight + bias). ggml_norm normalizes on ne[0]
|
|
// which is `hidden` here, matching PyTorch nn.LayerNorm(hidden).
|
|
struct ggml_tensor * ln1 = ggml_norm(ctx, x, tr->norm_eps);
|
|
ln1 = ggml_mul(ctx, ln1, layer.input_norm_w);
|
|
ln1 = ggml_add(ctx, ln1, layer.input_norm_b);
|
|
|
|
struct ggml_tensor * q = ggml_mul_mat(ctx, layer.q_proj_w, ln1);
|
|
struct ggml_tensor * k = ggml_mul_mat(ctx, layer.k_proj_w, ln1);
|
|
struct ggml_tensor * v = ggml_mul_mat(ctx, layer.v_proj_w, ln1);
|
|
|
|
q = ggml_reshape_3d(ctx, q, hd, n_q_heads, T);
|
|
k = ggml_reshape_3d(ctx, k, hd, n_kv, T);
|
|
v = ggml_reshape_3d(ctx, v, hd, n_kv, T);
|
|
|
|
q = ggml_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tr->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
|
|
0.0f);
|
|
k = ggml_rope_ext(ctx, k, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tr->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f,
|
|
0.0f);
|
|
|
|
struct ggml_tensor * q_p = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));
|
|
struct ggml_tensor * k_p = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));
|
|
struct ggml_tensor * v_p = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));
|
|
|
|
struct ggml_tensor * scores = ggml_mul_mat(ctx, k_p, q_p);
|
|
float scale = 1.0f / sqrtf((float) hd);
|
|
scores = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f);
|
|
|
|
struct ggml_tensor * attn = ggml_mul_mat(ctx, v_p, scores);
|
|
attn = ggml_cont(ctx, ggml_permute(ctx, attn, 0, 2, 1, 3));
|
|
attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T);
|
|
|
|
struct ggml_tensor * o = ggml_mul_mat(ctx, layer.o_proj_w, attn);
|
|
|
|
o = ggml_mul(ctx, o, layer.attn_scale);
|
|
x = ggml_add(ctx, x, o);
|
|
|
|
// MLP block: pre-LayerNorm (with bias) + fc1 -> GELU(erf) -> fc2 + LayerScale + residual.
|
|
struct ggml_tensor * ln2 = ggml_norm(ctx, x, tr->norm_eps);
|
|
ln2 = ggml_mul(ctx, ln2, layer.post_attn_norm_w);
|
|
ln2 = ggml_add(ctx, ln2, layer.post_attn_norm_b);
|
|
|
|
struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.fc1_w, ln2);
|
|
mlp = ggml_gelu_erf(ctx, mlp);
|
|
mlp = ggml_mul_mat(ctx, layer.fc2_w, mlp);
|
|
|
|
mlp = ggml_mul(ctx, mlp, layer.mlp_scale);
|
|
x = ggml_add(ctx, x, mlp);
|
|
|
|
(void) hidden;
|
|
return x;
|
|
}
|
|
|
|
// Full encoder transformer forward. No top-level input_proj or output_proj
|
|
// brackets: the SEANet output already has hidden_size channels.
|
|
// x : [hidden, T] f32 C-first
|
|
// positions: [T] i32
|
|
// mask : [T, T] f32 additive
|
|
// Returns [hidden, T] f32 C-first.
|
|
static struct ggml_tensor * enc_trans_forward(struct ggml_context * ctx,
|
|
const QwenEncoderTransformer * tr,
|
|
struct ggml_tensor * x,
|
|
struct ggml_tensor * positions,
|
|
struct ggml_tensor * mask) {
|
|
int T = (int) x->ne[1];
|
|
for (int l = 0; l < tr->num_layers; l++) {
|
|
x = enc_trans_layer_forward(ctx, tr, tr->layers[l], x, positions, mask, T);
|
|
}
|
|
return x;
|
|
}
|