encoder: replicate pad on downsample, drop orphan sliding_window field

This commit is contained in:
Pascal
2026-05-11 05:43:06 +02:00
parent 7e01c71c58
commit 9264b737a9
4 changed files with 46 additions and 16 deletions
+31 -3
View File
@@ -130,20 +130,31 @@ static struct ggml_tensor * qwen_causal_trans_conv1d(struct ggml_context * ctx,
// w: [k, IC, OC] f32, source layout (K, IC, OC) maps to ggml ne directly
// b: [OC] f32 or NULL
// x: [T, IC] f32 T-first
// pad_mode: QWEN_PAD_CONSTANT (zero pad, default for SEANet and the DAC
// decoder) or QWEN_PAD_REPLICATE (edge pad, replicates the
// first / last frame to match Mimi's downsample which is the
// only conv passing pad_mode="replicate" upstream).
// Returns [ceil(T / stride), OC] f32 T-first.
enum QwenPadMode {
QWEN_PAD_CONSTANT = 0,
QWEN_PAD_REPLICATE = 1,
};
static struct ggml_tensor * qwen_causal_conv1d(struct ggml_context * ctx,
struct ggml_tensor * w,
struct ggml_tensor * b,
struct ggml_tensor * x,
int k,
int d,
int s = 1) {
int s = 1,
int pad_mode = QWEN_PAD_CONSTANT) {
int OC = (int) w->ne[2];
int kernel_eff = (k - 1) * d + 1;
int padding_tot = kernel_eff - s;
// Mimi extra padding: ensures the causal conv lands on a stride boundary
// by adding zeros on the right before the convolution.
// by extending the input on the right with zeros or replicated edges
// depending on pad_mode.
int T = (int) x->ne[0];
int n_frames = (T + padding_tot - kernel_eff + s - 1) / s + 1;
int ideal_len = (n_frames - 1) * s + kernel_eff - padding_tot;
@@ -153,7 +164,24 @@ static struct ggml_tensor * qwen_causal_conv1d(struct ggml_context * ctx,
}
struct ggml_tensor * y = x;
if (padding_tot > 0 || extra_pad > 0) {
if (pad_mode == QWEN_PAD_REPLICATE) {
// Edge pad: repeat x[t=0] padding_tot times on the left and x[t=T-1]
// extra_pad times on the right via a single ggml_repeat per side.
int IC = (int) x->ne[1];
if (padding_tot > 0) {
struct ggml_tensor * first = ggml_view_2d(ctx, x, 1, IC, x->nb[1], 0);
struct ggml_tensor * tmpl = ggml_new_tensor_2d(ctx, x->type, padding_tot, IC);
struct ggml_tensor * lp = ggml_repeat(ctx, first, tmpl);
y = ggml_concat(ctx, lp, y, 0);
}
if (extra_pad > 0) {
size_t last_off = (size_t) (T - 1) * x->nb[1];
struct ggml_tensor * last = ggml_view_2d(ctx, x, 1, IC, x->nb[1], last_off);
struct ggml_tensor * tmpl = ggml_new_tensor_2d(ctx, x->type, extra_pad, IC);
struct ggml_tensor * rp = ggml_repeat(ctx, last, tmpl);
y = ggml_concat(ctx, y, rp, 0);
}
} else if (padding_tot > 0 || extra_pad > 0) {
y = ggml_pad_ext(ctx, y, padding_tot, extra_pad, 0, 0, 0, 0, 0, 0);
}
+5 -2
View File
@@ -57,11 +57,14 @@ static void qwen_encoder_downsample_free(QwenEncoderDownsample * d) {
}
}
// Forward: causal Conv1d k=4 stride=2, no bias.
// Forward: causal Conv1d k=4 stride=2, no bias. The Mimi downsample is
// the only conv on this side that ships with pad_mode="replicate"
// hardcoded upstream (transformers MimiModel.__init__), unlike SEANet
// and the encoder transformer which inherit config.pad_mode='constant'.
// x: [T, 512] f32 T-first
// Returns [ceil(T/2), 512] f32 T-first.
static struct ggml_tensor * qwen_encoder_downsample_forward(struct ggml_context * ctx,
const QwenEncoderDownsample * d,
struct ggml_tensor * x) {
return qwen_causal_conv1d(ctx, d->weight, NULL, x, d->kernel, 1, d->stride);
return qwen_causal_conv1d(ctx, d->weight, NULL, x, d->kernel, 1, d->stride, QWEN_PAD_REPLICATE);
}
+8 -10
View File
@@ -8,7 +8,8 @@
// - No biases on q/k/v/o projections
// - 8 attention heads instead of 16
// - intermediate_size 2048 instead of 1024
// - Causal sliding window 250 frames instead of 72
// - 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
//
@@ -59,7 +60,6 @@ struct QwenEncoderTransformer {
int num_kv_heads;
int head_dim;
int intermediate_size;
int sliding_window;
float rope_theta;
float norm_eps;
@@ -76,7 +76,6 @@ static bool qwen_encoder_transformer_load(QwenEncoderTransformer * tr, const GGU
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->sliding_window = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.sliding_window");
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");
@@ -119,9 +118,9 @@ static bool qwen_encoder_transformer_load(QwenEncoderTransformer * tr, const GGU
fprintf(stderr,
"[EncTransformer] Loaded: %d layers, hidden %d, heads %d/%d, head_dim %d, "
"FFN %d, RoPE theta %.0f, sliding window %d\n",
"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, tr->sliding_window);
tr->intermediate_size, tr->rope_theta);
return true;
}
@@ -137,11 +136,10 @@ static void qwen_encoder_transformer_free(QwenEncoderTransformer * tr) {
}
// Build a [T, T] additive causal mask (0 where allowed, -inf where masked).
// Pure causal : k <= q. Even though the upstream config carries a
// sliding_window field, MimiAttention's eager forward does not apply it
// (only the eager attention_mask is used) and MimiTransformerModel calls
// create_causal_mask() which is non-sliding. The Qwen3TTS encoder side
// inherits this convention, so we mirror it bit for bit here.
// 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 qwen_encoder_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++) {
+2 -1
View File
@@ -9,7 +9,8 @@
//
// Phase 4.1+ : eager attention, full F32 compute, no KV cache. The
// graph is built from scratch at every call ; a generation loop will
// later wrap this with a sliding KV cache.
// later wrap this with a plain causal KV cache (the talker Python
// reference uses pure causal attention, no sliding window).
//
// Optional dump_dir captures bisect-layer activations and the final
// logits in the same f32 binary format the Python reference produces,