diff --git a/src/causal-trans-conv.h b/src/causal-trans-conv.h index 53c6afd..9f4ba01 100644 --- a/src/causal-trans-conv.h +++ b/src/causal-trans-conv.h @@ -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); } diff --git a/src/encoder-downsample.h b/src/encoder-downsample.h index b51a53d..1930e04 100644 --- a/src/encoder-downsample.h +++ b/src/encoder-downsample.h @@ -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); } diff --git a/src/encoder-transformer.h b/src/encoder-transformer.h index 7305d0b..36a9e5a 100644 --- a/src/encoder-transformer.h +++ b/src/encoder-transformer.h @@ -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 & dst) { dst.assign((size_t) T * (size_t) T, -INFINITY); for (int q = 0; q < T; q++) { diff --git a/src/talker-forward.h b/src/talker-forward.h index 81b7d53..dda6336 100644 --- a/src/talker-forward.h +++ b/src/talker-forward.h @@ -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,