encoder: replicate pad on downsample, drop orphan sliding_window field
This commit is contained in:
+31
-3
@@ -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);
|
||||
}
|
||||
|
||||
|
||||
@@ -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,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++) {
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user