From 317ef255c441a7e304b2ac96d4705ace01a60656 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Sun, 13 Sep 2026 12:13:15 +0800 Subject: [PATCH 1/2] slim down parakeet memory usage --- src/arch/parakeet/encoder.cpp | 18 ++-- src/conformer/conformer.cpp | 164 ++++++++++++++++++---------------- src/conformer/conformer.h | 18 ++-- 3 files changed, 105 insertions(+), 95 deletions(-) diff --git a/src/arch/parakeet/encoder.cpp b/src/arch/parakeet/encoder.cpp index a26af118..06cd5059 100644 --- a/src/arch/parakeet/encoder.cpp +++ b/src/arch/parakeet/encoder.cpp @@ -325,14 +325,15 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } const bool var_len_masks = batch_var_len && n_batch > 1; conf::ConvPolicy policy{}; - policy.direct_pw = conf::detect_direct_pw(backend_name); - policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); - policy.direct_dw_in_pre_encode = detect_direct_dw_in_pre_encode(backend_name); + policy.direct_pw = conf::detect_direct_pw(backend_name); + policy.direct_conv0_in_pre_encode = true; + policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); + policy.direct_dw_in_pre_encode = detect_direct_dw_in_pre_encode(backend_name); // Cache-aware streaming (NeMo causal_downsampling=true) uses // CausalConv2D for the pre-encode subsample (left=k-1, right=stride-1). // Inferred from the attention style — only ChunkedLimited is causal. // Independent of the conformer conv-module's conv_context. - policy.causal_pre_encode = (hp.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited); + policy.causal_pre_encode = (hp.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited); EncoderBuild eb{}; @@ -702,11 +703,12 @@ EncoderBuild build_encoder_graph_streaming(ggml_context * ctx, const char * backend_name, bool spk_supervision) { conf::ConvPolicy policy{}; - policy.direct_pw = conf::detect_direct_pw(backend_name); - policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); - policy.direct_dw_in_pre_encode = detect_direct_dw_in_pre_encode(backend_name); + policy.direct_pw = conf::detect_direct_pw(backend_name); + policy.direct_conv0_in_pre_encode = true; + policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); + policy.direct_dw_in_pre_encode = detect_direct_dw_in_pre_encode(backend_name); // Causal pre-encode is the cache-aware streaming convention only. - policy.causal_pre_encode = (hp.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited); + policy.causal_pre_encode = (hp.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited); EncoderBuild eb{}; diff --git a/src/conformer/conformer.cpp b/src/conformer/conformer.cpp index 3ecfe27d..c1ddacda 100644 --- a/src/conformer/conformer.cpp +++ b/src/conformer/conformer.cpp @@ -623,88 +623,91 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, p = pos_proj; } - // Position mask / bias: matrix_bd = rel_shift(q_v @ p^T), truncated. - ggml_tensor * matrix_bd = ggml_mul_mat(ctx, p, q_v); - - // Local-attention pad/slice. The standard rel_shift trick assumes - // matrix_bd has shape [2T_q-1, T_q]: row r corresponds to relative - // offset (T_q-1-r). For local attention pos_emb is shorter - // ([W_left+W_right+1]) where row r corresponds to offset (W_left-r). - // Bring matrix_bd back to the [2T_q-1, T_q] shape by: - // - prepending (T_q-1-W_left) rows of -INF (or slicing them off - // when the audio is so short the window already covers it), - // - appending (T_q-1-W_right) rows of -INF (or slicing). - // After this, rel_shift + the existing T_q×T_q view land each - // out-of-window position at -INF, which softmax zeroes out. With - // both window sides == -1 (full attention) this block is skipped. - if (is_local) { - const int top_pad = static_cast(T_q) - 1 - W_left; - if (top_pad > 0) { - ggml_tensor * top_template = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, top_pad, T_q, n_head, B); - ggml_tensor * top = ggml_fill(ctx, top_template, -INFINITY); - matrix_bd = ggml_concat(ctx, top, matrix_bd, /*dim=*/0); - } else if (top_pad < 0) { - const int kept = static_cast(matrix_bd->ne[0]) + top_pad; - matrix_bd = ggml_view_4d(ctx, matrix_bd, kept, T_q, n_head, B, matrix_bd->nb[1], matrix_bd->nb[2], - matrix_bd->nb[3], (-top_pad) * matrix_bd->nb[0]); - matrix_bd = ggml_cont(ctx, matrix_bd); + auto shifted_view = [&](ggml_tensor * scores, int64_t heads) { + // The relative shift followed by the T_kv slice maps + // out[k, q] = in[k - q + T_q - 1, q]. Advancing q therefore + // moves one source row forward and one position backward. + return ggml_view_4d(ctx, scores, T_kv, T_q, heads, B, scores->nb[1] - scores->nb[0], scores->nb[2], + scores->nb[3], /*offset=*/(T_q - 1) * scores->nb[0]); + }; + + // A full-attention flash mask has an independent relative-position + // matmul for each head. Build and narrow those one head at a time so + // the allocator never needs the [2*T-1, T, H] F32 result at once. + // Masked/local cases retain the common path below. + const bool split_flash_mask = flash && !is_local && !is_chunked && params.attn_pad_mask == nullptr; + ggml_tensor * matrix_bd = nullptr; + if (split_flash_mask) { + for (int h = 0; h < n_head; ++h) { + ggml_tensor * p_h = ggml_view_4d(ctx, p, head_dim, pos_len, 1, 1, p->nb[1], p->nb[2], p->nb[3], + static_cast(h) * p->nb[2]); + ggml_tensor * q_v_h = ggml_view_4d(ctx, q_v, head_dim, T_q, 1, B, q_v->nb[1], q_v->nb[2], q_v->nb[3], + static_cast(h) * q_v->nb[2]); + ggml_tensor * head_mask = ggml_mul_mat(ctx, p_h, q_v_h); + head_mask = shifted_view(head_mask, 1); + head_mask = ggml_cont(ctx, head_mask); + head_mask = ggml_scale(ctx, head_mask, scale); + head_mask = ggml_cast(ctx, head_mask, GGML_TYPE_F16); + matrix_bd = matrix_bd == nullptr ? head_mask : ggml_concat(ctx, matrix_bd, head_mask, /*dim=*/2); } - const int bot_pad = static_cast(T_q) - 1 - W_right; - if (bot_pad > 0) { - ggml_tensor * bot_template = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, bot_pad, T_q, n_head, B); - ggml_tensor * bot = ggml_fill(ctx, bot_template, -INFINITY); - matrix_bd = ggml_concat(ctx, matrix_bd, bot, /*dim=*/0); - } else if (bot_pad < 0) { - const int kept = static_cast(matrix_bd->ne[0]) + bot_pad; - matrix_bd = ggml_view_4d(ctx, matrix_bd, kept, T_q, n_head, B, matrix_bd->nb[1], matrix_bd->nb[2], - matrix_bd->nb[3], /*offset=*/0); - matrix_bd = ggml_cont(ctx, matrix_bd); + } else { + // Position mask / bias: matrix_bd = rel_shift(q_v @ p^T), truncated. + matrix_bd = ggml_mul_mat(ctx, p, q_v); + + // Local-attention pad/slice. The standard rel_shift trick assumes + // matrix_bd has shape [2T_q-1, T_q]: row r corresponds to relative + // offset (T_q-1-r). For local attention pos_emb is shorter + // ([W_left+W_right+1]) where row r corresponds to offset (W_left-r). + // Bring matrix_bd back to the [2T_q-1, T_q] shape by padding or + // slicing both ends. Out-of-window positions then land at -INF. + if (is_local) { + const int top_pad = static_cast(T_q) - 1 - W_left; + if (top_pad > 0) { + ggml_tensor * top_template = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, top_pad, T_q, n_head, B); + ggml_tensor * top = ggml_fill(ctx, top_template, -INFINITY); + matrix_bd = ggml_concat(ctx, top, matrix_bd, /*dim=*/0); + } else if (top_pad < 0) { + const int kept = static_cast(matrix_bd->ne[0]) + top_pad; + matrix_bd = ggml_view_4d(ctx, matrix_bd, kept, T_q, n_head, B, matrix_bd->nb[1], matrix_bd->nb[2], + matrix_bd->nb[3], (-top_pad) * matrix_bd->nb[0]); + matrix_bd = ggml_cont(ctx, matrix_bd); + } + const int bot_pad = static_cast(T_q) - 1 - W_right; + if (bot_pad > 0) { + ggml_tensor * bot_template = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, bot_pad, T_q, n_head, B); + ggml_tensor * bot = ggml_fill(ctx, bot_template, -INFINITY); + matrix_bd = ggml_concat(ctx, matrix_bd, bot, /*dim=*/0); + } else if (bot_pad < 0) { + const int kept = static_cast(matrix_bd->ne[0]) + bot_pad; + matrix_bd = ggml_view_4d(ctx, matrix_bd, kept, T_q, n_head, B, matrix_bd->nb[1], matrix_bd->nb[2], + matrix_bd->nb[3], /*offset=*/0); + matrix_bd = ggml_cont(ctx, matrix_bd); + } } - } - - // rel_shift generalizes to the rectangular case: with input - // [T_q + T_kv - 1, T_q] it yields out[k, q] = in[k - q + T_q - 1, q], - // i.e. row k holds the score of key k against query q for relative - // offset (T_kv - 1) - (k - q + T_q - 1) = (T_kv - T_q) + q - k — - // exactly the query-at-absolute-position (T_kv - T_q + q) semantics - // the streaming x_q path needs. The square offline case is the - // T_q == T_kv specialization. The zero column injected by the trick - // only lands at k >= T_kv, which the view below slices off. - matrix_bd = rel_shift(ctx, matrix_bd); - matrix_bd = ggml_view_4d(ctx, matrix_bd, T_kv, T_q, n_head, B, matrix_bd->nb[1], matrix_bd->nb[2], matrix_bd->nb[3], - /*offset=*/0); - // The view is non-contiguous (nb[1] stays at parent's - // pos_len*es), but it IS contiguous-rows. The flash path calls - // ggml_scale which wants full contiguity, so cont there. The - // manual path only feeds matrix_bd into ggml_add(kq, matrix_bd), - // which handles contiguous-rows inputs on both CPU and Metal. - if (flash) { - matrix_bd = ggml_cont(ctx, matrix_bd); - } - // ChunkedLimited mask. Caller provided a [T_q, T_q, 1, 1] F32 - // tensor with 0 on allowed (q, k) pairs and -INF outside the - // [q_chunk - left_chunks, q_chunk] band. Broadcasts across n_head. - // -INF and 0 are scale-invariant so this can be added before the - // pre-scale that the flash path applies below. - if (is_chunked && params.attn_chunked_mask != nullptr) { - matrix_bd = ggml_add(ctx, matrix_bd, params.attn_chunked_mask); - } + matrix_bd = shifted_view(matrix_bd, n_head); + // The manual path accepts contiguous rows. Flash needs a fully + // contiguous mask before scaling and narrowing it below. + if (flash) { + matrix_bd = ggml_cont(ctx, matrix_bd); + } - // Variable-length batch key-padding mask. [T_k, 1, 1, B] additive - // (-INF on padded keys) broadcasts over queries and heads. Added here - // so it applies on both the flash and manual paths (matrix_bd is the - // flash mask and the manual additive bias alike). -INF / 0 are - // scale-invariant, so adding before the flash pre-scale is fine. - if (params.attn_pad_mask != nullptr) { - matrix_bd = ggml_add(ctx, matrix_bd, params.attn_pad_mask); + // Chunked and key-padding masks broadcast over heads and queries. + if (is_chunked && params.attn_chunked_mask != nullptr) { + matrix_bd = ggml_add(ctx, matrix_bd, params.attn_chunked_mask); + } + if (params.attn_pad_mask != nullptr) { + matrix_bd = ggml_add(ctx, matrix_bd, params.attn_pad_mask); + } } ggml_tensor * o; if (flash) { - matrix_bd = ggml_scale(ctx, matrix_bd, scale); - matrix_bd = ggml_cast(ctx, matrix_bd, GGML_TYPE_F16); + if (!split_flash_mask) { + matrix_bd = ggml_scale(ctx, matrix_bd, scale); + matrix_bd = ggml_cast(ctx, matrix_bd, GGML_TYPE_F16); + } // Optionally cast K/V activations to a narrower type to // reduce bandwidth in the attention kernel. GGML_TYPE_COUNT @@ -1006,10 +1009,17 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, // conv0 (standard 2D conv: 1 in, channels out, k=3 s=2) x = pad_causal(x); - x = ggml_conv_2d(ctx, pe.conv0_w, x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); + if (policy.direct_conv0_in_pre_encode) { + x = ggml_conv_2d_direct(ctx, pe.conv0_w, x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } else { + x = ggml_conv_2d(ctx, pe.conv0_w, x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } x = add_conv_bias(ctx, x, pe.conv0_b); x = name_prefixed(x, name_prefix, "conv0"); x = ggml_relu(ctx, x); diff --git a/src/conformer/conformer.h b/src/conformer/conformer.h index cb1c547a..7a823798 100644 --- a/src/conformer/conformer.h +++ b/src/conformer/conformer.h @@ -102,17 +102,15 @@ struct BlockView { }; // Per-family conv dispatch policy. direct_pw is shared (detect_direct_pw is -// the same for every family today); direct_dw splits between the block site -// (direct_dw_in_block: the conformer block's 1-D depthwise after GLU) and the -// pre_encode site (direct_dw_in_pre_encode: the stride-2 2-D depthwise), since -// the two have different shapes and per-family backend choices. Defaults are -// conservative: direct_pw true (pointwise is direct mul_mat everywhere), both -// direct_dw_* false (im2col), which is safe on Metal where the direct 2-D -// depthwise kernel is not implemented for all shapes. +// the same for every family today). The pre-encode conv0 and depthwise sites +// have separate direct-op controls because their shapes and backend tradeoffs +// differ from the block convolutions. Defaults keep the established im2col +// paths except for direct_pw; families opt into direct pre-encode ops. struct ConvPolicy { - bool direct_pw = true; - bool direct_dw_in_block = false; - bool direct_dw_in_pre_encode = false; + bool direct_pw = true; + bool direct_conv0_in_pre_encode = false; + bool direct_dw_in_block = false; + bool direct_dw_in_pre_encode = false; // Causal pre_encode convolutions. NeMo's cache-aware streaming swaps // every Conv2d in ConvSubsampling for CausalConv2D, padding From f567229432ac19faa0671520b1437ee0aa957e6c Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Sun, 13 Sep 2026 13:42:09 +0800 Subject: [PATCH 2/2] stale comment --- src/conformer/conformer.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/conformer/conformer.cpp b/src/conformer/conformer.cpp index c1ddacda..5be61d22 100644 --- a/src/conformer/conformer.cpp +++ b/src/conformer/conformer.cpp @@ -651,7 +651,7 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, matrix_bd = matrix_bd == nullptr ? head_mask : ggml_concat(ctx, matrix_bd, head_mask, /*dim=*/2); } } else { - // Position mask / bias: matrix_bd = rel_shift(q_v @ p^T), truncated. + // Compute position scores; shifted_view narrows and shifts them below. matrix_bd = ggml_mul_mat(ctx, p, q_v); // Local-attention pad/slice. The standard rel_shift trick assumes