Skip to content

slim down parakeet memory usage #161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cjpais merged 2 commits into from
Sep 13, 2026
+105 −95
Merged

slim down parakeet memory usage #161

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
317ef25
slim down parakeet memory usage
cjpais Sep 13, 2026
f567229
stale comment
cjpais Sep 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 10 additions & 8 deletions src/arch/parakeet/encoder.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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{};

Expand Down Expand Up @@ -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{};

Expand Down
164 changes: 87 additions & 77 deletions src/conformer/conformer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>(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<int>(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<size_t>(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<size_t>(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<int>(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<int>(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 {
// 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
// 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<int>(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<int>(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<int>(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<int>(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
Expand Down Expand Up @@ -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);
Expand Down
18 changes: 8 additions & 10 deletions src/conformer/conformer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading