From 707aa904924e299c982857346f1277a5ca58f75e Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Tue, 20 Oct 2026 23:54:13 -0700 Subject: [PATCH 1/6] backend: log a graph's largest tensors at debug (TRANSCRIBE_GRAPH_TOP) The sched-scratch line from #217 says how big a graph's committed buffer is; it does not say which tensors make it that big. Finding that out meant bisecting a graph by hand. TRANSCRIBE_GRAPH_TOP=N logs the N largest nodes of every graph allocated through alloc_sched_graph: op, type, shape, assigned backend and name, sorted by size, plus the total if nothing were reused (an upper bound; the real buffer is the live set). Views are skipped since they own no memory. Unset or non-positive is inert, so this is a debugging aid with no behaviour change. It is what attributed both memory branches that follow: granite_nar's [100352, T_enc] BPE CTC head and qwen3_asr's per-chunk conv activations were invisible in the scratch totals alone. --- src/transcribe-backend.cpp | 44 ++++++++++++++++++++++++++++++++++++++ src/transcribe-backend.h | 3 +++ 2 files changed, 47 insertions(+) diff --git a/src/transcribe-backend.cpp b/src/transcribe-backend.cpp index 41bd8071..cf20cd31 100644 --- a/src/transcribe-backend.cpp +++ b/src/transcribe-backend.cpp @@ -7,8 +7,10 @@ #include "transcribe-backend.h" #include "ggml.h" +#include "transcribe-env.h" #include "transcribe-log.h" +#include #include #include #include @@ -201,6 +203,47 @@ void log_sched_scratch(ggml_backend_sched_t sched, const char * tag, int line, l } } +// TRANSCRIBE_GRAPH_TOP=N: log the N largest nodes of every graph allocated +// through alloc_sched_graph, with op, type, shape and assigned backend. The +// scratch line above says how big a graph is; this says which tensors make it +// that big. Sizes are per tensor, not the live set, so a tensor that is freed +// before the next big one is allocated still shows up. Debugging aid only. +void log_graph_top(ggml_backend_sched_t sched, struct ggml_cgraph * graph, const char * tag, int line) { + const char * env_n = env::str("TRANSCRIBE_GRAPH_TOP"); + if (env_n == nullptr) { + return; + } + const int top_n = std::atoi(env_n); + if (top_n <= 0) { + return; + } + std::vector nodes; + const int n_nodes = ggml_graph_n_nodes(graph); + nodes.reserve(static_cast(n_nodes)); + size_t total = 0; + for (int i = 0; i < n_nodes; ++i) { + ggml_tensor * t = ggml_graph_node(graph, i); + if (t->view_src != nullptr) { + continue; // views own no memory + } + total += ggml_nbytes(t); + nodes.push_back(t); + } + const size_t keep = std::min(nodes.size(), static_cast(top_n)); + std::partial_sort(nodes.begin(), nodes.begin() + static_cast(keep), nodes.end(), + [](const ggml_tensor * a, const ggml_tensor * b) { return ggml_nbytes(a) > ggml_nbytes(b); }); + log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, "%s@%d: graph top %zu of %d nodes (%.1f MB if nothing were reused)", tag, line, + keep, n_nodes, static_cast(total) / 1e6); + for (size_t i = 0; i < keep; ++i) { + const ggml_tensor * t = nodes[i]; + ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, const_cast(t)); + log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, " %9.1f MB %-14s %-5s [%lld, %lld, %lld, %lld] %s %s", + static_cast(ggml_nbytes(t)) / 1e6, ggml_op_desc(t), ggml_type_name(t->type), + static_cast(t->ne[0]), static_cast(t->ne[1]), static_cast(t->ne[2]), + static_cast(t->ne[3]), backend != nullptr ? ggml_backend_name(backend) : "-", t->name); + } +} + } // namespace void safe_backend_free(ggml_backend_t backend) noexcept { @@ -244,6 +287,7 @@ bool alloc_sched_graph(ggml_backend_sched_t sched, return false; } log_sched_scratch(sched, tag, line, extent); + log_graph_top(sched, graph, tag, line); return true; } diff --git a/src/transcribe-backend.h b/src/transcribe-backend.h index f1c106ec..82efa359 100644 --- a/src/transcribe-backend.h +++ b/src/transcribe-backend.h @@ -112,6 +112,9 @@ void release_compute_scratch(ggml_backend_sched_t & sched, struct ggml_context * // buffer's high-water mark, so per-token step reallocations stay silent. // `tag`/`line` identify the call site; `extent` (>= 0) is an optional // graph-length field to log alongside it (e.g. the encoder frame count). +// +// TRANSCRIBE_GRAPH_TOP=N additionally logs the N largest nodes of each graph +// (op, type, shape, backend, name), to attribute the scratch to tensors. bool alloc_sched_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph, const char * tag, From ea5df41eba6447bb11d345ce72d6d34dc49e97bc Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Tue, 20 Oct 2026 23:54:13 -0700 Subject: [PATCH 2/6] qwen3_asr: bound the encoder graph, and turn encoder flash on The encoder was the whole memory story: 6.0 GB of committed Metal scratch for an 811 MB model at 5 minutes, 8.0 GB RSS on CPU, growing ~linearly with audio length. Two causes, both measured with TRANSCRIBE_GRAPH_TOP: 1. The 3x Conv2d subsampler ran every chunk in one batch. Per chunk that is 6.1 MB of conv0 activation and 6.9 MB of conv1 im2col, so ~1.2 GB of scratch per minute of audio. Run it in groups of 16 chunks inside the one graph instead: each group's intermediates retire before the next group's are allocated, and only the [d_model, T_per_chunk, g] outputs are concatenated. Per-chunk results do not depend on the group size, and enc.subsample.out / enc.pos_add.out are bit-identical to before. Shared by the single-shot and batched graphs (subsample_chunks). 2. Encoder flash was off (encoder_use_flash = false since 1d244e5c, a Vulkan perf commit with no recorded reason), so every block materialised [T_enc, T_enc, n_heads] f32 scores: T^2, 0.9 GB at 5 minutes and 13.6 GB at 20. Cause 1 hid it at 5 minutes. Flash is GQA/bidirectional-native here and needs no mask, which also drops the all-zero [T_enc, T_enc] mask that was built host-side and uploaded per run. Metal, Q8_0, dots-full: scratch 6051 -> 397 MB at 306 s, and 1224 s now completes at 1797 MB (previously not runnable). CPU RSS at 306 s 8035 -> 2281 MB. Encode 2843 -> 2120 ms on Metal; CPU encode unchanged (1283 -> 1284 ms at 60 s), so the flash default is safe on both. Transcripts are byte-identical to main at 60/120/306 s, and LibriSpeech test-clean Q8_0, b8, 2620 utts is 2.13% WER against the catalog's published 2.11% [1.95, 2.27]. TRANSCRIBE_NO_FLASH still selects the old attention path. --- src/arch/qwen3_asr/encoder.cpp | 200 +++++++++++++++++++-------------- src/arch/qwen3_asr/model.cpp | 22 ++-- 2 files changed, 124 insertions(+), 98 deletions(-) diff --git a/src/arch/qwen3_asr/encoder.cpp b/src/arch/qwen3_asr/encoder.cpp index c569d56b..8bb9c6a2 100644 --- a/src/arch/qwen3_asr/encoder.cpp +++ b/src/arch/qwen3_asr/encoder.cpp @@ -9,6 +9,7 @@ #include "transcribe-debug.h" #include "transcribe-log.h" +#include #include #include #include @@ -70,6 +71,10 @@ std::vector build_sinusoid_pe(int32_t d_model, int32_t length, double max return pe; } +// The single-shot graph no longer takes this mask: it is all zeros, which +// is exactly soft_max_ext / flash_attn_ext with a null mask, and as a dense +// [T_enc, T_enc] input it was a T^2 host + device buffer (~1 GB each at +// 20 minutes). Kept as the documented statement of the attention pattern. std::vector build_cu_seqlens_mask(const EncoderTiming & t, const QwenAsrHParams & hp) { (void) hp; const int32_t T = t.T_enc; @@ -125,9 +130,108 @@ ggml_tensor * linear(ggml_context * ctx, ggml_tensor * x, ggml_tensor * w, ggml_ return y; } -// One encoder block: pre-LN self-attention (bidirectional, full- -// sequence mask supplied by the caller) + pre-LN GELU FFN. Residuals -// are full 1.0. +// Conv chunks per subsample group. The 3x Conv2d frontend is independent per +// n_window*2-mel chunk, but running every chunk in one batch makes its +// activations linear in audio length with a huge constant: conv0's output is +// [50, 64, 480] f32 (6.1 MB) per chunk and conv1's im2col another 6.9 MB, so +// ~1.2 GB of scheduler scratch per minute of audio. Grouping bounds that to a +// fixed size: each group's intermediates die before the next group's are +// allocated, and only the small [d_model, T_per_chunk, g] outputs are +// concatenated. The reference does the same thing with conv_chunksize (500); +// ours is smaller because the scheduler commits the whole group at once. +// Per-chunk results do not depend on the group size. +constexpr int64_t kConvChunkGroup = 16; + +// Conv subsample + conv_out (no PE) for `mel` with ne=[mel_per_chunk, +// n_mels, 1, n] (any n). Returns [d_model, T_per_chunk, n], or nullptr if the +// conv output width does not match T_per_chunk. +ggml_tensor * subsample_group(ggml_context * ctx, + const QwenAsrWeights & weights, + int64_t ds_h, + int64_t T_per_chunk, + ggml_tensor * mel) { + const int64_t n = mel->ne[3]; + + // Reference layout: PyTorch Conv2d on [B, 1, H=n_mels, W=mel_per_chunk]. + // For ggml_conv_2d we use [W=mel_per_chunk, H=n_mels, C=1, N=B] — + // ggml swaps (W, H) relative to PyTorch, but the 3x3 kernel is + // symmetric and stride/pad are (2,2)/(1,1), so the arithmetic is + // invariant. Kernels are stored as ne=[KW=3, KH=3, IC, OC]. + ggml_tensor * x = mel; + x = ggml_conv_2d(ctx, weights.enc_subsample.conv0_w, x, + /*s0=*/2, /*s1=*/2, /*p0=*/1, /*p1=*/1, /*d0=*/1, /*d1=*/1); + x = add_conv_bias(ctx, x, weights.enc_subsample.conv0_b); + x = ggml_gelu_erf(ctx, x); + x = ggml_conv_2d(ctx, weights.enc_subsample.conv1_w, x, 2, 2, 1, 1, 1, 1); + x = add_conv_bias(ctx, x, weights.enc_subsample.conv1_b); + x = ggml_gelu_erf(ctx, x); + x = ggml_conv_2d(ctx, weights.enc_subsample.conv2_w, x, 2, 2, 1, 1, 1, 1); + x = add_conv_bias(ctx, x, weights.enc_subsample.conv2_b); + x = ggml_gelu_erf(ctx, x); + // Now ne = [W=mel_ds, H=n_mels_ds, C=ds_h, N=n]. + + const int64_t W_ds = x->ne[0]; + const int64_t H_ds = x->ne[1]; + if (W_ds != T_per_chunk) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "qwen3_asr encoder: post-conv W=%lld does not match " + "per_chunk_aftercnn=%lld", + static_cast(W_ds), static_cast(T_per_chunk)); + return nullptr; + } + + // Reshape for the linear projection. Reference does: + // permute(0, 3, 1, 2).contiguous().view(b, t, c*f) + // i.e. [B, C, F_ds, T_ds] -> [B, T_ds, C, F_ds] -> [B, T_ds, C*F_ds]. + // The flat inner axis is (c * F_ds + f). + // + // In ggml ne layout we have [W=T_ds, H=F_ds, C=ds_h, N=B]. To get a + // flat axis of size ds_h * F_ds where c is the slower sub-axis and + // f the faster — matching the reference — we want axes in order + // [f, c, T, B]: ne = [F_ds, ds_h, T_ds, B]. ggml_permute uses INVERSE + // semantics vs PyTorch: the i-th argument says which NEW axis old axis + // i goes to (new[a_i] = old[i]). To mirror PyTorch's `permute(0,3,1,2)` + // on [B, C, F, T] (equivalently, map [W=13, H=16, C=480, N=11] → + // [F=16, C=480, T=13, B=11]), we need new[0]=old[1], new[1]=old[2], + // new[2]=old[0], new[3]=old[3], which is ggml_permute args (2, 0, 1, 3). + x = ggml_permute(ctx, x, /*a0=*/2, /*a1=*/0, /*a2=*/1, /*a3=*/3); + x = ggml_cont(ctx, x); + x = ggml_reshape_3d(ctx, x, H_ds * ds_h, T_per_chunk, n); + + // Linear conv_out: [ds_h*F_ds, d_model] weight; maps to d_model. + x = ggml_mul_mat(ctx, weights.enc_subsample.conv_out, x); + return x; // ne = [d_model, T_per_chunk, n] +} + +// subsample_group over every chunk of `mel` (ne[3] = n_chunks), in groups of +// kConvChunkGroup, concatenated back along the chunk axis. Returns +// [d_model, T_per_chunk, n_chunks], or nullptr on a shape mismatch. +ggml_tensor * subsample_chunks(ggml_context * ctx, + const QwenAsrWeights & weights, + int64_t ds_h, + int64_t T_per_chunk, + ggml_tensor * mel) { + const int64_t n_chunks = mel->ne[3]; + ggml_tensor * out = nullptr; + for (int64_t c0 = 0; c0 < n_chunks; c0 += kConvChunkGroup) { + const int64_t n = std::min(kConvChunkGroup, n_chunks - c0); + ggml_tensor * g = (n == n_chunks) ? mel : + ggml_view_4d(ctx, mel, mel->ne[0], mel->ne[1], mel->ne[2], n, mel->nb[1], + mel->nb[2], mel->nb[3], static_cast(c0) * mel->nb[3]); + ggml_tensor * y = subsample_group(ctx, weights, ds_h, T_per_chunk, g); + if (y == nullptr) { + return nullptr; + } + // Chaining the concat keeps the groups in dependency order, so the + // allocator retires one group's activations before the next group's. + out = (out == nullptr) ? y : ggml_concat(ctx, out, y, /*dim=*/2); + } + return out; +} + +// One encoder block: pre-LN self-attention (bidirectional; `mask` may be +// null for unmasked full attention) + pre-LN GELU FFN. Residuals are +// full 1.0. ggml_tensor * build_enc_block(ggml_context * ctx, ggml_tensor * x, ggml_tensor * mask, @@ -243,63 +347,14 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, named(eb.pos_emb_in, "enc.pos_emb.in"); ggml_set_input(eb.pos_emb_in); - eb.mask_in = ggml_new_tensor_2d(ctx, use_flash ? GGML_TYPE_F16 : GGML_TYPE_F32, T_enc, T_enc); - named(eb.mask_in, "enc.attn_mask.in"); - ggml_set_input(eb.mask_in); - - // ----- Subsample: 3x Conv2d + GELU + conv_out linear ----- - // Reference layout: PyTorch Conv2d on [B, 1, H=n_mels, W=mel_per_chunk]. - // For ggml_conv_2d we use [W=mel_per_chunk, H=n_mels, C=1, N=B] — - // ggml swaps (W, H) relative to PyTorch, but the 3x3 kernel is - // symmetric and stride/pad are (2,2)/(1,1), so the arithmetic is - // invariant. Kernels are stored as ne=[KW=3, KH=3, IC, OC]. - ggml_tensor * x = eb.mel_in; - - x = ggml_conv_2d(ctx, weights.enc_subsample.conv0_w, x, - /*s0=*/2, /*s1=*/2, /*p0=*/1, /*p1=*/1, /*d0=*/1, /*d1=*/1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv0_b); - x = ggml_gelu_erf(ctx, x); - - x = ggml_conv_2d(ctx, weights.enc_subsample.conv1_w, x, 2, 2, 1, 1, 1, 1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv1_b); - x = ggml_gelu_erf(ctx, x); - - x = ggml_conv_2d(ctx, weights.enc_subsample.conv2_w, x, 2, 2, 1, 1, 1, 1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv2_b); - x = ggml_gelu_erf(ctx, x); - // Now ne = [W=mel_ds, H=n_mels_ds, C=ds_h, N=n_chunks]. + // No attention mask: full unmasked attention (see build_cu_seqlens_mask). + eb.mask_in = nullptr; - const int64_t W_ds = x->ne[0]; - const int64_t H_ds = x->ne[1]; - if (W_ds != T_per_chunk) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "qwen3_asr encoder: post-conv W=%lld does not match " - "per_chunk_aftercnn=%lld", - static_cast(W_ds), static_cast(T_per_chunk)); + // ----- Subsample: 3x Conv2d + GELU + conv_out linear, per chunk group ----- + ggml_tensor * x = subsample_chunks(ctx, weights, ds_h, T_per_chunk, eb.mel_in); + if (x == nullptr) { return eb; } - - // Reshape for the linear projection. Reference does: - // permute(0, 3, 1, 2).contiguous().view(b, t, c*f) - // i.e. [B, C, F_ds, T_ds] -> [B, T_ds, C, F_ds] -> [B, T_ds, C*F_ds]. - // The flat inner axis is (c * F_ds + f). - // - // In ggml ne layout we have [W=T_ds, H=F_ds, C=ds_h, N=B]. To get a - // flat axis of size ds_h * F_ds where c is the slower sub-axis and - // f the faster — matching the reference — we want axes in order - // [f, c, T, B]: ne = [F_ds, ds_h, T_ds, B]. That's a permute(1, 2, 0, 3). - // ggml_permute uses INVERSE semantics vs PyTorch: the i-th argument - // says which NEW axis old axis i goes to (new[a_i] = old[i]). To - // mirror PyTorch's `permute(0,3,1,2)` on [B, C, F, T] (equivalently, - // map [W=13, H=16, C=480, N=11] → [F=16, C=480, T=13, B=11]), we - // need new[0]=old[1], new[1]=old[2], new[2]=old[0], new[3]=old[3], - // which is ggml_permute args (2, 0, 1, 3). - x = ggml_permute(ctx, x, /*a0=*/2, /*a1=*/0, /*a2=*/1, /*a3=*/3); - x = ggml_cont(ctx, x); - x = ggml_reshape_3d(ctx, x, H_ds * ds_h, T_per_chunk, n_chunks); - - // Linear conv_out: [ds_h*F_ds, d_model] weight; maps to d_model. - x = ggml_mul_mat(ctx, weights.enc_subsample.conv_out, x); // ne = [d_model, T_per_chunk, n_chunks] named(x, "enc.subsample.out"); eb.dumps.subsample_out = x; @@ -346,7 +401,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // ----- 18 encoder blocks ----- const int n_layers = static_cast(weights.enc_blocks.size()); for (int i = 0; i < n_layers; ++i) { - x = build_enc_block(ctx, x, eb.mask_in, weights.enc_blocks[i], static_cast(d_model), + x = build_enc_block(ctx, x, /*mask=*/nullptr, weights.enc_blocks[i], static_cast(d_model), static_cast(n_heads), use_flash); if (i == 0) { named(x, "enc.block.0.out"); @@ -454,35 +509,12 @@ EncoderBuildBatched build_encoder_graph_batched(ggml_context * ctx, named(eb.mask_in, "enc.attn_mask.in"); ggml_set_input(eb.mask_in); - // ----- Subsample: 3x Conv2d + GELU (per-chunk over N = B*n_chunks_max) ----- - ggml_tensor * x = eb.mel_in; - x = ggml_conv_2d(ctx, weights.enc_subsample.conv0_w, x, 2, 2, 1, 1, 1, 1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv0_b); - x = ggml_gelu_erf(ctx, x); - x = ggml_conv_2d(ctx, weights.enc_subsample.conv1_w, x, 2, 2, 1, 1, 1, 1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv1_b); - x = ggml_gelu_erf(ctx, x); - x = ggml_conv_2d(ctx, weights.enc_subsample.conv2_w, x, 2, 2, 1, 1, 1, 1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv2_b); - x = ggml_gelu_erf(ctx, x); - // ne = [W=T_per_chunk, H=n_mels_ds, C=ds_h, N]. - - const int64_t W_ds = x->ne[0]; - const int64_t H_ds = x->ne[1]; - if (W_ds != T_per_chunk) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "qwen3_asr encoder(batched): post-conv W=%lld != " - "per_chunk_aftercnn=%lld", - static_cast(W_ds), static_cast(T_per_chunk)); + // ----- Subsample: 3x Conv2d + GELU + conv_out (per-chunk over N = B*n_chunks_max) ----- + ggml_tensor * x = subsample_chunks(ctx, weights, ds_h, T_per_chunk, eb.mel_in); // [d_model, T_per_chunk, N] + if (x == nullptr) { return eb; } - // [W=T_ds, H=F_ds, C=ds_h, N] -> [F_ds, ds_h, T_ds, N] (see single-shot). - x = ggml_permute(ctx, x, 2, 0, 1, 3); - x = ggml_cont(ctx, x); - x = ggml_reshape_3d(ctx, x, H_ds * ds_h, T_per_chunk, N); - x = ggml_mul_mat(ctx, weights.enc_subsample.conv_out, x); // [d_model, T_per_chunk, N] - // Add sinusoidal PE (broadcast across the N = B*n_chunks_max axis). x = ggml_add(ctx, x, eb.pos_emb_in); diff --git a/src/arch/qwen3_asr/model.cpp b/src/arch/qwen3_asr/model.cpp index 66a7868b..81ca130e 100644 --- a/src/arch/qwen3_asr/model.cpp +++ b/src/arch/qwen3_asr/model.cpp @@ -290,7 +290,12 @@ transcribe_status init_context(transcribe_model * model, cc->kv_type = params->kv_type; cc->n_ctx = transcribe_session_params_n_ctx(params); - cc->encoder_use_flash = false; + // Encoder flash is what keeps long audio bounded: without it every block + // materialises [T_enc, T_enc, n_heads] f32 scores (~0.9 GB at 5 minutes, + // ~13.6 GB at 20). Measured on M4 Metal at 306 s it is also faster + // (encode 2843 -> 2120 ms) with an identical transcript; CPU is neutral. + // TRANSCRIBE_NO_FLASH still turns it off. + cc->encoder_use_flash = true; cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); @@ -693,19 +698,8 @@ transcribe_status run(transcribe_session * session, ggml_backend_tensor_set(eb.pos_emb_in, pe.data(), 0, pe.size() * sizeof(float)); } - // Attention mask (block-diagonal from cu_seqlens). - { - std::vector mask = build_cu_seqlens_mask(timing, cm->hparams); - if (cc->encoder_use_flash) { - std::vector mask_f16(mask.size()); - for (size_t i = 0; i < mask.size(); ++i) { - mask_f16[i] = ggml_fp32_to_fp16(mask[i]); - } - ggml_backend_tensor_set(eb.mask_in, mask_f16.data(), 0, mask_f16.size() * sizeof(ggml_fp16_t)); - } else { - ggml_backend_tensor_set(eb.mask_in, mask.data(), 0, mask.size() * sizeof(float)); - } - } + // No attention mask: the encoder attends unmasked over the whole + // utterance (see build_cu_seqlens_mask), so there is nothing to upload. transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); From f539a0f9fd179905d5a142e5fac4abdd7630cb5c Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Tue, 20 Oct 2026 23:54:13 -0700 Subject: [PATCH 3/6] granite_nar: bound the encoder and decoder graphs The encoder committed 6.5 GB of scratch at 306 s for a 2.4 GB model and the run then rejected the input anyway; at 240 s it was 11.1 GB RSS. Three causes, measured with TRANSCRIBE_GRAPH_TOP: 1. The BPE CTC head ran on every frame: [100352, T_enc] f32, i.e. 6.1 GB at 306 s (50 encoder frames per second), then copied to the host whole. The reference pools it over windows of 4 frames with importance weights that sum to 1, and a head is affine: sum_t w_t (W h_t + b) = W (sum_t w_t h_t) + b so pool the 1024-wide hidden states first, then run the head on n_windows = T_enc / 4 columns. Tiles of 256 windows bound the logits to ~103 MB, and ggml_argmax runs in-graph, so only one id per window (plus each window's importance mass, for the all-blank rule) reaches the host. Same arithmetic, 4x fewer FLOPs in the head. collapse_bpe_window_ids keeps compute_bpe_ctc_initial_hypothesis's exact semantics, including both blank-id schemes. 2. The depthwise conv ran through conv_1d_dw_f32's im2col + mul_mat, which inflates the activation by conv_kernel (15) to feed a matmul whose K is that same kernel: [15, T, 2048] f32, 1.9 GB at 306 s. Use conv_2d_dw_direct_f32, as granite5_ctc already does. 3. The decoder built [T_total, T_total, 16] f32 scores (539 MB at 180 s, 1.07 GB at the 4096-token cap) plus the K/V ggml_repeat copies, and wrote the whole [vocab, n_text] logits tensor back to the host. Flash attention is GQA-native and needs no mask here (bidirectional); the LM head runs in tiles of 256 text rows with the argmax in-graph. The full logits form is still built when tensor dumping is on, since dec.text_logits is a parity tap. Decode 5.20 -> 4.16 s at 180 s. Metal, Q8_0: 180 s 8369 -> 1032 MB RSS, 8849 -> 1039 MB footprint, scratch 3832 -> 708 MB; 240 s 11146 -> 1339 MB RSS. 306 s still rejects (3059 audio + 1893 text tokens > 4096), which is the model's context, not memory. Transcripts identical to main at 60/120/180/240 s and 0/647 differ on fleurs-en; LibriSpeech test-clean Q8_0, b8, 2620 utts is 1.28% WER against the catalog's published 1.29%. Branch-vs-main tensor drift enters at enc.block.0.post_conv (2.4e-4, the direct-depthwise accumulation order) and reaches enc.block.15.out 0.25 max / 0.019 mean; the reference tolerances are Q8_0-limited on main too (2/15 either way). --- src/arch/granite_nar/decoder.cpp | 149 ++++++++++++++++++----------- src/arch/granite_nar/decoder.h | 26 ++++- src/arch/granite_nar/encoder.cpp | 125 ++++++++++++++++++++---- src/arch/granite_nar/encoder.h | 44 ++++++--- src/arch/granite_nar/granite_nar.h | 8 +- src/arch/granite_nar/model.cpp | 100 ++++++++++--------- 6 files changed, 304 insertions(+), 148 deletions(-) diff --git a/src/arch/granite_nar/decoder.cpp b/src/arch/granite_nar/decoder.cpp index 80f2c4a9..6e08cb5e 100644 --- a/src/arch/granite_nar/decoder.cpp +++ b/src/arch/granite_nar/decoder.cpp @@ -64,7 +64,8 @@ ggml_tensor * block_bidi(ggml_context * ctx, const GraniteNarDecBlock & view, const BlockParams & params, int T_total, - ggml_tensor * positions) { + ggml_tensor * positions, + bool use_flash) { const int64_t n_heads = params.n_heads; const int64_t n_kv_heads = params.n_kv_heads; const int64_t n_groups = n_heads / n_kv_heads; @@ -95,34 +96,45 @@ ggml_tensor * block_bidi(ggml_context * ctx, ggml_tensor * K_att = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); ggml_tensor * V_att = ggml_cont(ctx, ggml_permute(ctx, V, 0, 2, 1, 3)); - // GQA: repeat each KV head across the n_groups Q heads. - ggml_tensor * K_full; - ggml_tensor * V_full; - if (n_groups == 1) { - K_full = K_att; - V_full = V_att; + ggml_tensor * o = nullptr; + if (use_flash) { + // flash_attn_ext broadcasts the n_kv_heads K/V across the n_heads Q + // heads itself and never materialises [T_total, T_total, n_heads] + // scores (~1.07 GB f32 per layer at the 4096-token cap). No mask: + // bidirectional. Result is [head_dim, n_heads, T_total]. + o = ggml_flash_attn_ext(ctx, Q_att, K_att, V_att, /*mask=*/nullptr, scale_attn, /*max_bias=*/0.0f, + /*logit_softcap=*/0.0f); + o = ggml_reshape_2d(ctx, o, static_cast(n_heads) * head_dim, T_total); } else { - ggml_tensor * K_4d = ggml_reshape_4d(ctx, K_att, head_dim, T_total, 1, n_kv_heads); - ggml_tensor * V_4d = ggml_reshape_4d(ctx, V_att, head_dim, T_total, 1, n_kv_heads); - ggml_tensor * K_tpl = ggml_new_tensor_4d(ctx, K_att->type, head_dim, T_total, n_groups, n_kv_heads); - ggml_tensor * V_tpl = ggml_new_tensor_4d(ctx, V_att->type, head_dim, T_total, n_groups, n_kv_heads); - ggml_tensor * K_rep = ggml_repeat(ctx, K_4d, K_tpl); - ggml_tensor * V_rep = ggml_repeat(ctx, V_4d, V_tpl); - K_full = ggml_reshape_3d(ctx, K_rep, head_dim, T_total, n_heads); - V_full = ggml_reshape_3d(ctx, V_rep, head_dim, T_total, n_heads); - } - - // scores = (K^T Q) * scale_attn. ggml_mul_mat(K, Q) yields - // [T_total_k, T_total_q, n_heads]. - ggml_tensor * kq = ggml_mul_mat(ctx, K_full, Q_att); - // No mask (bidirectional). Use soft_max_ext with mask=nullptr for the - // scaling, then softmax. - ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, /*mask=*/nullptr, scale_attn, /*max_bias=*/0.0f); + // GQA: repeat each KV head across the n_groups Q heads. + ggml_tensor * K_full; + ggml_tensor * V_full; + if (n_groups == 1) { + K_full = K_att; + V_full = V_att; + } else { + ggml_tensor * K_4d = ggml_reshape_4d(ctx, K_att, head_dim, T_total, 1, n_kv_heads); + ggml_tensor * V_4d = ggml_reshape_4d(ctx, V_att, head_dim, T_total, 1, n_kv_heads); + ggml_tensor * K_tpl = ggml_new_tensor_4d(ctx, K_att->type, head_dim, T_total, n_groups, n_kv_heads); + ggml_tensor * V_tpl = ggml_new_tensor_4d(ctx, V_att->type, head_dim, T_total, n_groups, n_kv_heads); + ggml_tensor * K_rep = ggml_repeat(ctx, K_4d, K_tpl); + ggml_tensor * V_rep = ggml_repeat(ctx, V_4d, V_tpl); + K_full = ggml_reshape_3d(ctx, K_rep, head_dim, T_total, n_heads); + V_full = ggml_reshape_3d(ctx, V_rep, head_dim, T_total, n_heads); + } - ggml_tensor * V_t = ggml_cont(ctx, ggml_permute(ctx, V_full, 1, 0, 2, 3)); - ggml_tensor * o = ggml_mul_mat(ctx, V_t, kq_soft); - o = ggml_cont(ctx, ggml_permute(ctx, o, 0, 2, 1, 3)); - o = ggml_reshape_2d(ctx, o, static_cast(n_heads) * head_dim, T_total); + // scores = (K^T Q) * scale_attn. ggml_mul_mat(K, Q) yields + // [T_total_k, T_total_q, n_heads]. + ggml_tensor * kq = ggml_mul_mat(ctx, K_full, Q_att); + // No mask (bidirectional). Use soft_max_ext with mask=nullptr for the + // scaling, then softmax. + ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, /*mask=*/nullptr, scale_attn, /*max_bias=*/0.0f); + + ggml_tensor * V_t = ggml_cont(ctx, ggml_permute(ctx, V_full, 1, 0, 2, 3)); + o = ggml_mul_mat(ctx, V_t, kq_soft); + o = ggml_cont(ctx, ggml_permute(ctx, o, 0, 2, 1, 3)); + o = ggml_reshape_2d(ctx, o, static_cast(n_heads) * head_dim, T_total); + } o = ggml_mul_mat(ctx, view.attn_o_w, o); o = ggml_scale(ctx, o, params.residual_mul); @@ -145,7 +157,9 @@ ForwardBuild build_forward_graph(ggml_context * ctx, const GraniteNarWeights & weights, const GraniteNarHParams & hp, int n_audio_tokens, - int n_text) { + int n_text, + bool use_flash, + bool full_logits) { ForwardBuild fb{}; fb.n_audio_tokens = n_audio_tokens; fb.n_text = n_text; @@ -208,7 +222,7 @@ ForwardBuild build_forward_graph(ggml_context * ctx, // Block stack. for (int il = 0; il < n_layer; ++il) { - x = block_bidi(ctx, x, weights.dec_blocks[il], params, fb.T_total, fb.positions_in); + x = block_bidi(ctx, x, weights.dec_blocks[il], params, fb.T_total, fb.positions_in, use_flash); } // Final RMSNorm. @@ -220,22 +234,39 @@ ForwardBuild build_forward_graph(ggml_context * ctx, ggml_element_size(x) * hidden * static_cast(n_audio_tokens)); text_x = ggml_cont(ctx, text_x); - // LM head (tied to token_embd). The Granite-4 head divides by - // logits_scaling; we mirror it so dec.text_logits aligns numerically with - // the reference dump. Argmax is scale-invariant, so the transcript is - // unchanged by this division. - ggml_tensor * logits = ggml_mul_mat(ctx, weights.dec_embed.token_w, text_x); - if (hp.dec_logits_scaling > 0.0f && hp.dec_logits_scaling != 1.0f) { - logits = ggml_scale(ctx, logits, 1.0f / hp.dec_logits_scaling); + if (full_logits) { + // LM head (tied to token_embd). The Granite-4 head divides by + // logits_scaling; we mirror it so dec.text_logits aligns numerically + // with the reference dump. Argmax is scale-invariant, so the + // transcript is unchanged by this division. + ggml_tensor * logits = ggml_mul_mat(ctx, weights.dec_embed.token_w, text_x); + if (hp.dec_logits_scaling > 0.0f && hp.dec_logits_scaling != 1.0f) { + logits = ggml_scale(ctx, logits, 1.0f / hp.dec_logits_scaling); + } + // ne = [vocab, n_text] + named(logits, "dec.text_logits"); + fb.out = logits; + fb.dumps.text_logits = logits; + ggml_set_output(logits); + transcribe::debug::mark_tensor_for_dump(logits); + ggml_build_forward_expand(gf, fb.out); + } else { + // Argmax form: [vocab, n_text] f32 is ~0.8 GB at a 2000-row + // hypothesis, plus its host copy, for a result that is one id per row. + // Tiles are expanded in order so each tile's logits are retired before + // the next tile's are allocated. logits_scaling is skipped: dividing + // by a positive constant does not move the argmax. + for (int r0 = 0; r0 < n_text; r0 += kTextTileRows) { + const int n = std::min(kTextTileRows, n_text - r0); + ggml_tensor * tv = + ggml_view_2d(ctx, text_x, hidden, n, text_x->nb[1], static_cast(r0) * text_x->nb[1]); + ggml_tensor * id = ggml_argmax(ctx, ggml_mul_mat(ctx, weights.dec_embed.token_w, tv)); // [n] i32 + ggml_set_output(id); + fb.text_ids.push_back(id); + ggml_build_forward_expand(gf, id); + } } - // ne = [vocab, n_text] - named(logits, "dec.text_logits"); - fb.out = logits; - fb.dumps.text_logits = logits; - ggml_set_output(logits); - transcribe::debug::mark_tensor_for_dump(logits); - - ggml_build_forward_expand(gf, fb.out); + if (fb.dumps.flat_embeds) { ggml_build_forward_expand(gf, fb.dumps.flat_embeds); } @@ -256,13 +287,28 @@ void add_insertion_slots(const std::vector & hyp_ids, int32_t eos_id, s } } +void collapse_drop_eos(const std::vector & row_ids, int32_t eos_id, std::vector & out_ids) { + out_ids.clear(); + int32_t prev = -1; + for (const int32_t best : row_ids) { + if (best == eos_id) { + prev = best; + continue; + } + if (best == prev) { + continue; + } + out_ids.push_back(best); + prev = best; + } +} + void argmax_collapse_drop_eos(const std::vector & text_logits, int vocab, int n_text, int32_t eos_id, std::vector & out_ids) { - out_ids.clear(); - int32_t prev = -1; + std::vector row_ids(static_cast(std::max(n_text, 0))); for (int i = 0; i < n_text; ++i) { const float * row = text_logits.data() + static_cast(i) * vocab; int32_t best = 0; @@ -273,16 +319,9 @@ void argmax_collapse_drop_eos(const std::vector & text_logits, best = v; } } - if (best == eos_id) { - prev = best; - continue; - } - if (best == prev) { - continue; - } - out_ids.push_back(best); - prev = best; + row_ids[static_cast(i)] = best; } + collapse_drop_eos(row_ids, eos_id, out_ids); } } // namespace transcribe::granite_nar diff --git a/src/arch/granite_nar/decoder.h b/src/arch/granite_nar/decoder.h index 174baa14..238553a3 100644 --- a/src/arch/granite_nar/decoder.h +++ b/src/arch/granite_nar/decoder.h @@ -30,6 +30,7 @@ #include "weights.h" #include +#include struct ggml_context; struct ggml_cgraph; @@ -50,8 +51,11 @@ struct ForwardBuild { ggml_tensor * text_ids_in = nullptr; // [n_text] i32 ggml_tensor * positions_in = nullptr; // [T_total] i32 RoPE positions - // Output (lm_head sliced over the text portion). - ggml_tensor * out = nullptr; // [vocab, n_text] + // Output (lm_head sliced over the text portion). Exactly one of the two + // forms is built (see build_forward_graph's full_logits): + ggml_tensor * out = nullptr; // [vocab, n_text] full logits + std::vector text_ids; // argmax per text row, in tiles + // of kTextTileRows, each [n] i32 DecoderDumps dumps{}; ggml_cgraph * graph = nullptr; @@ -61,13 +65,26 @@ struct ForwardBuild { int T_total = 0; }; +// Text rows per LM-head tile in the argmax form: bounds the +// [vocab, rows] f32 logits to ~103 MB (100352 x 256). +constexpr int kTextTileRows = 256; + // Build the bidirectional editor forward graph. n_text is the length of // the post-insertion-slots text sequence (2 * n_hyp + 1). +// +// use_flash: attention through flash_attn_ext (GQA-native, no +// [T_total, T_total, n_heads] scores) instead of mul_mat + soft_max_ext. +// full_logits: emit the whole [vocab, n_text] `dec.text_logits` (needed for +// tensor dumps). Otherwise the LM head runs in tiles of kTextTileRows with +// the argmax in-graph and only ids reach the host; the transcript is the +// same because argmax is invariant to the positive logits_scaling. ForwardBuild build_forward_graph(ggml_context * ctx, const GraniteNarWeights & weights, const GraniteNarHParams & hp, int n_audio_tokens, - int n_text); + int n_text, + bool use_flash, + bool full_logits); // Host-side helper: insert eos_id between every adjacent pair (and at // the head/tail) of `hyp_ids`. Length grows from n to 2n + 1. @@ -82,4 +99,7 @@ void argmax_collapse_drop_eos(const std::vector & text_logits, int32_t eos_id, std::vector & out_ids); +// The collapse half of argmax_collapse_drop_eos, over per-row argmax ids. +void collapse_drop_eos(const std::vector & row_ids, int32_t eos_id, std::vector & out_ids); + } // namespace transcribe::granite_nar diff --git a/src/arch/granite_nar/encoder.cpp b/src/arch/granite_nar/encoder.cpp index 57a3a516..4fc3c456 100644 --- a/src/arch/granite_nar/encoder.cpp +++ b/src/arch/granite_nar/encoder.cpp @@ -4,10 +4,10 @@ // self-attention, GLU conv module with conv_expansion=2, macaron FFN, // mid-layer self-conditioned CTC bypass). NLE additions: // -// - A second CTC head over a BPE vocab. We emit -// frame-level logits as `enc.ctc_bpe_logits` here; the -// posterior-weighted window pool + greedy decode runs host-side at -// run() time. +// - A second CTC head over a BPE vocab, applied to posterior-weighted +// window pools of the final hidden state. The pool and the argmax run +// in-graph; only per-window ids reach the host, where the greedy +// collapse runs at run() time. // - All-hidden-states capture: we tap the per-block POST-LN output at // indices specified by hp.enc_layer_indices (e.g. [4, 8, 12, -1] // 1-indexed → block outputs after 3, 7, 11, 15). These are @@ -190,11 +190,24 @@ ggml_tensor * conv_module(ggml_context * ctx, } x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); const int padding = (conv_kernel - 1) / 2; - x = transcribe::conformer::conv_1d_dw_f32(ctx, b.conv_depthwise_w, x, - /*stride=*/1, /*padding=*/padding, /*dilation=*/1); - x = transcribe::conformer::fused_batch_norm(ctx, x, bn_fused_scale, bn_fused_bias); - x = ggml_silu(ctx, x); - x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); + // Direct depthwise op, not conv_1d_dw_f32's im2col + mul_mat: im2col + // materialises a [conv_kernel, T, inner_dim] f32 copy (15 x 2048 x 50 + // frames/s: ~1.9 GB at 5 minutes, the encoder's largest tensor once the + // BPE head is pooled) to feed a matmul whose K is only conv_kernel. Same + // swap as granite5_ctc (see its encoder.cpp); not bitwise identical + // (accumulation order), so it is gated by the family's dump parity. + { + ggml_tensor * knl = ggml_reshape_4d(ctx, b.conv_depthwise_w, conv_kernel, 1, 1, inner_dim); + ggml_tensor * img = ggml_reshape_4d(ctx, x, T, 1, inner_dim, 1); + x = transcribe::conformer::conv_2d_dw_direct_f32(ctx, knl, img, + /*s0=*/1, /*s1=*/1, + /*p0=*/padding, /*p1=*/0, + /*d0=*/1, /*d1=*/1); + x = ggml_reshape_2d(ctx, x, T, inner_dim); + } + x = transcribe::conformer::fused_batch_norm(ctx, x, bn_fused_scale, bn_fused_bias); + x = ggml_silu(ctx, x); + x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); { ggml_tensor * pw2 = ggml_reshape_2d(ctx, b.conv_pointwise2_w, inner_dim, d_model); x = ggml_mul_mat(ctx, pw2, x); @@ -392,14 +405,56 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } } - // Frame-level BPE CTC head (the pool happens host-side). - ggml_tensor * ctc_bpe = nullptr; - if (weights.enc_top.ctc_bpe_w != nullptr) { - ctc_bpe = ggml_mul_mat(ctx, weights.enc_top.ctc_bpe_w, x); - ctc_bpe = ggml_add(ctx, ctc_bpe, weights.enc_top.ctc_bpe_b); - named(ctc_bpe, "enc.ctc_bpe_logits"); - eb.ctc_bpe_logits = ctc_bpe; - ggml_set_output(ctc_bpe); + // Pooled BPE CTC head. The reference forms, per pool window of P frames, + // the importance-weighted mean of the frame-level BPE logits and takes its + // argmax. Done literally that is a [bpe_output_dim, T_enc] f32 tensor + // (100352 x 50 frames/s: 6.1 GB of scratch at 5 minutes, then copied to + // host). The head is affine and the window weights sum to 1, so + // sum_t w_t (W h_t + b) = W (sum_t w_t h_t) + b + // and we pool the 1024-wide hidden states first, then run the head on + // n_windows = T_enc / P columns in tiles, with the argmax in-graph. Only + // the per-window ids and masses leave the device. + if (weights.enc_top.ctc_bpe_w != nullptr && eb.mid_blank_probs != nullptr && hp.enc_bpe_pool_window > 0) { + const int64_t P = hp.enc_bpe_pool_window; + const int64_t n_win = (T_enc + P - 1) / P; + const int64_t T_win = n_win * P; + + // importance = 1 - p_blank(mid CTC), zero-padded to whole windows so + // the partial last window pools only its real frames. + ggml_tensor * imp = ggml_scale_bias(ctx, eb.mid_blank_probs, -1.0f, 1.0f); // [T_enc] + ggml_tensor * h = x; // [d_model, T_enc] + if (T_win > T_enc) { + imp = ggml_pad(ctx, imp, static_cast(T_win - T_enc), 0, 0, 0); + h = ggml_pad(ctx, h, 0, static_cast(T_win - T_enc), 0, 0); + } + imp = ggml_reshape_2d(ctx, imp, P, n_win); + ggml_tensor * mass = ggml_sum_rows(ctx, imp); // [1, n_win] + named(mass, "enc.bpe_window_mass"); + ggml_set_output(mass); + eb.bpe_window_mass = mass; + eb.n_bpe_windows = static_cast(n_win); + + // All-blank windows (mass ~ 0) are blank on the host; the clamp only + // keeps their weights finite. + ggml_tensor * wt = ggml_div(ctx, imp, ggml_clamp(ctx, mass, 1e-9f, std::numeric_limits::max())); + + // pooled[:, w] = sum_p wt[p, w] * h[:, w*P + p] + ggml_tensor * hw = + ggml_mul(ctx, ggml_reshape_3d(ctx, h, d_model, P, n_win), ggml_reshape_3d(ctx, wt, 1, P, n_win)); + hw = ggml_cont(ctx, ggml_permute(ctx, hw, 1, 0, 2, 3)); // [P, d_model, n_win] + ggml_tensor * pooled = ggml_reshape_2d(ctx, ggml_sum_rows(ctx, hw), d_model, n_win); + named(pooled, "enc.bpe_pooled"); + + for (int64_t w0 = 0; w0 < n_win; w0 += kBpeTileWindows) { + const int64_t n = std::min(kBpeTileWindows, n_win - w0); + ggml_tensor * pv = + ggml_view_2d(ctx, pooled, d_model, n, pooled->nb[1], static_cast(w0) * pooled->nb[1]); + ggml_tensor * lg = ggml_mul_mat(ctx, weights.enc_top.ctc_bpe_w, pv); + lg = ggml_add(ctx, lg, weights.enc_top.ctc_bpe_b); + ggml_tensor * id = ggml_argmax(ctx, lg); // [n] i32 + ggml_set_output(id); + eb.bpe_window_ids.push_back(id); + } } // Channel concat of captured layer outputs. Walk captures in order @@ -428,12 +483,17 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } ggml_build_forward_expand(eb.graph, eb.cat_out); ggml_build_forward_expand(eb.graph, eb.ctc_logits); - if (eb.ctc_bpe_logits) { - ggml_build_forward_expand(eb.graph, eb.ctc_bpe_logits); - } if (eb.mid_blank_probs) { ggml_build_forward_expand(eb.graph, eb.mid_blank_probs); } + if (eb.bpe_window_mass) { + ggml_build_forward_expand(eb.graph, eb.bpe_window_mass); + } + // Expanded in order so each tile's logits are retired before the next + // tile's are allocated. + for (ggml_tensor * id : eb.bpe_window_ids) { + ggml_build_forward_expand(eb.graph, id); + } if (eb.dumps.input_linear_out) { ggml_build_forward_expand(eb.graph, eb.dumps.input_linear_out); } @@ -465,7 +525,30 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, return eb; } -// Host-side BPE CTC pool + greedy decode. +// Collapse for the in-graph pooled BPE head. + +void collapse_bpe_window_ids(const std::vector & window_ids, + const std::vector & window_mass, + int blank_id, + std::vector & out_token_ids) { + out_token_ids.clear(); + const size_t n_windows = std::min(window_ids.size(), window_mass.size()); + out_token_ids.reserve(n_windows); + // Same greedy + collapse-repeats + drop-blanks as + // compute_bpe_ctc_initial_hypothesis, including the all-blank-window rule + // and the two blank_id schemes. + const int shift = (blank_id == 0) ? 1 : 0; + int prev = -1; + for (size_t w = 0; w < n_windows; ++w) { + const int id = (window_mass[w] <= 1e-9f) ? blank_id : window_ids[w]; + if (id != blank_id && id != prev) { + out_token_ids.push_back(id - shift); + } + prev = id; + } +} + +// Host-side BPE CTC pool + greedy decode (reference formulation). void compute_bpe_ctc_initial_hypothesis(const std::vector & importance_non_blank, const std::vector & ctc_bpe_logits, diff --git a/src/arch/granite_nar/encoder.h b/src/arch/granite_nar/encoder.h index 5fc24338..ec2ba6fa 100644 --- a/src/arch/granite_nar/encoder.h +++ b/src/arch/granite_nar/encoder.h @@ -9,8 +9,9 @@ // window=4 pool of valid frames. We expose the bypass-step char-CTC // mid_logits (1024 → 348) as `enc.ctc_logits` — the exact tensor // the reference model computes at self_conditioning_layer for the -// self-conditioning residual. The pooled BPE head is computed -// host-side at run time. +// self-conditioning residual. The pooled BPE head runs in-graph on +// window-pooled hidden states and only per-window argmax ids leave the +// device (see build_encoder_graph). // 2. All-hidden-states capture: the projector consumes 4 encoder // hidden states (post-LN, pre-bypass at the chosen layer // boundaries; indices [4, 8, 12, -1] 1-indexed → layer outputs @@ -54,15 +55,19 @@ struct EncoderBuild { ggml_tensor * zero_pad = nullptr; // optional [hidden, T_pad-T_enc] // Outputs. - ggml_tensor * cat_out = nullptr; // [num_enc_layers * hidden, T_enc] - // the projector input - ggml_tensor * ctc_logits = nullptr; // [enc_out_dim=348, T_enc] - ggml_tensor * ctc_bpe_logits = nullptr; // [bpe_output_dim, T_enc] - // raw frame-level (no pool) - ggml_tensor * mid_blank_probs = nullptr; // [T_enc] — softmax(mid_ctc)[blank]. - // Used host-side as the BPE pool's - // importance weight (importance = - // 1 - blank_prob_mid). + ggml_tensor * cat_out = nullptr; // [num_enc_layers * hidden, T_enc] + // the projector input + ggml_tensor * ctc_logits = nullptr; // [enc_out_dim=348, T_enc] + ggml_tensor * mid_blank_probs = nullptr; // [T_enc] — softmax(mid_ctc)[blank]. + // importance = 1 - blank_prob_mid is + // the BPE pool's per-frame weight. + // Pooled BPE CTC head: argmax over bpe_output_dim per pool window, in + // consecutive tiles of windows (tile k covers windows + // [k * kBpeTileWindows, ...)), plus each window's total importance so the + // host can treat all-blank windows exactly like the reference. + std::vector bpe_window_ids; // each [n_tile] i32 + ggml_tensor * bpe_window_mass = nullptr; // [1, n_windows] f32 + int n_bpe_windows = 0; ggml_cgraph * graph = nullptr; @@ -96,7 +101,22 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, std::vector precompute_pos_rows(int context_size, int max_pos_emb); std::vector precompute_last_block_mask(int context_size, int t_enc_remainder); -// Host-side BPE CTC pool + greedy decode. +// Pool windows per BPE-head tile: bounds the [bpe_output_dim, tile] f32 +// logits the head materialises (~103 MB at 100352 x 256). +constexpr int kBpeTileWindows = 256; + +// Greedy collapse of the in-graph pooled BPE head: one argmax id per pool +// window and the window's total importance. Windows whose importance sum is +// <= 1e-9 emit blank, exactly as compute_bpe_ctc_initial_hypothesis does. +void collapse_bpe_window_ids(const std::vector & window_ids, + const std::vector & window_mass, + int blank_id, + std::vector & out_token_ids); + +// Host-side BPE CTC pool + greedy decode over frame-level logits. This is +// the reference formulation; run() uses the in-graph pooled head above +// (sum_t w_t (W h_t + b) = W (sum_t w_t h_t) + b because the window weights +// sum to 1) and keeps this as the documented semantics. // // Input: // importance_non_blank [T_enc] host row of (1 - blank_prob) values used diff --git a/src/arch/granite_nar/granite_nar.h b/src/arch/granite_nar/granite_nar.h index 500b8649..e33ba712 100644 --- a/src/arch/granite_nar/granite_nar.h +++ b/src/arch/granite_nar/granite_nar.h @@ -59,15 +59,13 @@ struct GraniteNarModel final : public transcribe_model { struct GraniteNarSession final : public transcribe_session { // Encoder output buffered between encode and projector/LM. std::vector mel_buf; - std::vector enc_cat_host; // [T_enc, num_encoder_layers * enc_hidden] - std::vector ctc_logits_host; // [T_enc, output_dim] - std::vector ctc_bpe_logits_host; // [N_valid, bpe_output_dim] flat - std::vector proj_out_host; // [n_audio_tokens, llm_dim] + std::vector enc_cat_host; // [T_enc, num_encoder_layers * enc_hidden] + std::vector proj_out_host; // [n_audio_tokens, llm_dim] int32_t t_enc = 0; int32_t n_audio_tokens = 0; bool encoder_use_flash = false; - bool decoder_use_flash = false; // bidirectional path doesn't use flash-attn KV cache + bool decoder_use_flash = true; // init_context sets the real default GraniteNarSession() = default; ~GraniteNarSession() override; diff --git a/src/arch/granite_nar/model.cpp b/src/arch/granite_nar/model.cpp index 0b02687a..63177f71 100644 --- a/src/arch/granite_nar/model.cpp +++ b/src/arch/granite_nar/model.cpp @@ -10,9 +10,9 @@ // forward pass — no KV cache, is_causal=False) // // The forward composes in one run() call as: -// PCM -> mel -> 2-frame stack -> encoder graph (emit cat_out, -// ctc_logits, ctc_bpe_logits) -> host CTC pool + greedy decode of -// the BPE head -> initial hypothesis -> add_insertion_slots -> +// PCM -> mel -> 2-frame stack -> encoder graph (emit cat_out, ctc_logits, +// in-graph pooled BPE head argmax per window) -> host greedy collapse +// -> initial hypothesis -> add_insertion_slots -> // text_ids -> decoder forward graph (with projector audio embeds // concatenated to the front of inputs_embeds) -> text_logits -> // argmax + collapse + drop EOS -> final transcript. @@ -366,8 +366,12 @@ transcribe_status init_context(transcribe_model * model, cc->n_threads = params->n_threads; cc->kv_type = params->kv_type; cc->n_ctx = transcribe_session_params_n_ctx(params); - cc->encoder_use_flash = false; - cc->decoder_use_flash = false; // bidirectional path doesn't use KV cache + cc->encoder_use_flash = false; // the Shaw block-local attention has no flash form + // The bidirectional editor attends over audio + text in one pass; + // flash keeps it from materialising [T_total, T_total, n_heads] scores + // (~1.07 GB f32 per layer at the 4096-token cap). TRANSCRIBE_NO_FLASH + // restores the mul_mat + soft_max path. + cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); *out_ctx = cc.release(); @@ -519,43 +523,21 @@ transcribe_status run(transcribe_session * ctx_base, cc->enc_cat_host.resize(static_cast(cat_h) * cat_T); ggml_backend_tensor_get(eb.cat_out, cc->enc_cat_host.data(), 0, cc->enc_cat_host.size() * sizeof(float)); - const int n_ctc_vocab = static_cast(eb.ctc_logits->ne[0]); - cc->ctc_logits_host.resize(static_cast(n_ctc_vocab) * cat_T); - ggml_backend_tensor_get(eb.ctc_logits, cc->ctc_logits_host.data(), 0, cc->ctc_logits_host.size() * sizeof(float)); - - int n_bpe_vocab = 0; - if (eb.ctc_bpe_logits != nullptr) { - n_bpe_vocab = static_cast(eb.ctc_bpe_logits->ne[0]); - cc->ctc_bpe_logits_host.resize(static_cast(n_bpe_vocab) * cat_T); - ggml_backend_tensor_get(eb.ctc_bpe_logits, cc->ctc_bpe_logits_host.data(), 0, - cc->ctc_bpe_logits_host.size() * sizeof(float)); - } - - std::vector mid_non_blank; - if (eb.mid_blank_probs != nullptr) { - std::vector mid_blank(t_enc); - ggml_backend_tensor_get(eb.mid_blank_probs, mid_blank.data(), 0, mid_blank.size() * sizeof(float)); - mid_non_blank.resize(t_enc); - for (int i = 0; i < t_enc; ++i) { - mid_non_blank[i] = 1.0f - mid_blank[i]; - } - } - - // ggml stores ne[0]=F as the innermost (fastest-varying) axis. For - // a tensor with ne=[F, T_enc], the host buffer layout in memory is - // buf[t * F + f] = element at ggml indices (f, t) - // which is the SAME as `[T_enc, F]` numpy row-major (T outer, F - // inner). No transpose needed: the host-side BPE pool reads - // ctc_bpe_logits_host[t * V + v] - // and gets the v-th logit at frame t correctly. - (void) n_ctc_vocab; - - // Initial BPE hypothesis. + // Initial BPE hypothesis from the in-graph pooled head (see + // build_encoder_graph): one argmax id per pool window, plus each window's + // importance mass for the all-blank rule. std::vector hyp_ids; - if (n_bpe_vocab > 0 && !mid_non_blank.empty()) { - compute_bpe_ctc_initial_hypothesis(mid_non_blank, cc->ctc_bpe_logits_host, n_bpe_vocab, t_enc, - cm->hparams.enc_bpe_pool_window, - /*blank_id=*/cm->hparams.enc_bpe_blank_id, hyp_ids); + if (eb.bpe_window_mass != nullptr && eb.n_bpe_windows > 0) { + std::vector window_mass(static_cast(eb.n_bpe_windows)); + ggml_backend_tensor_get(eb.bpe_window_mass, window_mass.data(), 0, window_mass.size() * sizeof(float)); + std::vector window_ids; + window_ids.reserve(static_cast(eb.n_bpe_windows)); + for (ggml_tensor * ids : eb.bpe_window_ids) { + const size_t off = window_ids.size(); + window_ids.resize(off + static_cast(ids->ne[0])); + ggml_backend_tensor_get(ids, window_ids.data() + off, 0, static_cast(ids->ne[0]) * sizeof(int32_t)); + } + collapse_bpe_window_ids(window_ids, window_mass, /*blank_id=*/cm->hparams.enc_bpe_blank_id, hyp_ids); } if (hyp_ids.empty()) { // No tokens — emit empty transcript and return. @@ -684,8 +666,12 @@ transcribe_status run(transcribe_session * ctx_base, return TRANSCRIBE_ERR_OOM; } } - ForwardBuild fb = build_forward_graph(dec_ctx, cm->weights, cm->hparams, cc->n_audio_tokens, n_text); - if (fb.graph == nullptr || fb.out == nullptr) { + // Full [vocab, n_text] logits only when tensors are being dumped (the + // dec.text_logits parity tap); otherwise argmax in-graph. + const bool full_logits = transcribe::debug::enabled(); + ForwardBuild fb = build_forward_graph(dec_ctx, cm->weights, cm->hparams, cc->n_audio_tokens, n_text, + cc->decoder_use_flash, full_logits); + if (fb.graph == nullptr || (fb.out == nullptr && fb.text_ids.empty())) { ggml_free(dec_ctx); return TRANSCRIBE_ERR_GGUF; } @@ -720,16 +706,26 @@ transcribe_status run(transcribe_session * ctx_base, try_dump("dec.flat_embeds", fb.dumps.flat_embeds, "decoder"); try_dump("dec.text_logits", fb.dumps.text_logits, "decoder"); - // Read text-portion logits. ggml stores ne[0]=vocab as the innermost - // axis, which means the host buffer layout is [t * vocab + v] — - // i.e. `[n_text, vocab]` numpy row-major already. No transpose. - const int64_t vocab = fb.out->ne[0]; - std::vector logits_host(static_cast(vocab) * n_text); - ggml_backend_tensor_get(fb.out, logits_host.data(), 0, logits_host.size() * sizeof(float)); - ggml_free(dec_ctx); - std::vector final_ids; - argmax_collapse_drop_eos(logits_host, static_cast(vocab), n_text, cm->hparams.dec_eos_id, final_ids); + if (fb.out != nullptr) { + // Read text-portion logits. ggml stores ne[0]=vocab as the innermost + // axis, which means the host buffer layout is [t * vocab + v] — + // i.e. `[n_text, vocab]` numpy row-major already. No transpose. + const int64_t vocab = fb.out->ne[0]; + std::vector logits_host(static_cast(vocab) * n_text); + ggml_backend_tensor_get(fb.out, logits_host.data(), 0, logits_host.size() * sizeof(float)); + argmax_collapse_drop_eos(logits_host, static_cast(vocab), n_text, cm->hparams.dec_eos_id, final_ids); + } else { + std::vector row_ids; + row_ids.reserve(static_cast(n_text)); + for (ggml_tensor * ids : fb.text_ids) { + const size_t off = row_ids.size(); + row_ids.resize(off + static_cast(ids->ne[0])); + ggml_backend_tensor_get(ids, row_ids.data() + off, 0, static_cast(ids->ne[0]) * sizeof(int32_t)); + } + collapse_drop_eos(row_ids, cm->hparams.dec_eos_id, final_ids); + } + ggml_free(dec_ctx); cc->full_text = cm->tok.decode(final_ids.data(), static_cast(final_ids.size())); cc->raw_text = cc->full_text; // no post-processing, via transcribe_raw_text From b490c12f9ca426b9e154009ab6402d4337518fac Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Wed, 21 Oct 2026 00:22:19 -0700 Subject: [PATCH 4/6] granite_nar, qwen3_asr: drop dead paths, harden the BPE window collapse Review follow-up on the two memory commits. - qwen3_asr::build_cu_seqlens_mask has no caller: the single-shot encoder passes mask=nullptr. An unused implementation can drift from the real path, so delete it and move its reasoning (the measured eager reference attends unmasked; vLLM's flash-attn-2 windowing scores worse) into the header's shape note, which also stops claiming mask_in is [T_enc, T_enc] when it is now always null. - granite_nar::compute_bpe_ctc_initial_hypothesis still has no production caller either, but it is the reference formulation the pooled in-graph head claims to reproduce, so it earns a test instead of a comment: tests/granite_nar_bpe_pool_unit.cpp builds the same affine head both ways (frame-level logits + host pool vs importance-weighted hidden pool + head + argmax) and requires identical tokens, including the all-blank-window rule, the partial tail window and both blank-id schemes. It also pins the collapse rules directly. - collapse_bpe_window_ids no longer truncates on a length mismatch: the ids and masses both come from tensors sized eb.n_bpe_windows, so a mismatch is a graph/caller bug. It now returns false and run() fails with TRANSCRIBE_ERR_BACKEND and a message naming both lengths, instead of emitting a silently shorter hypothesis. - tests/tolerances/granite_nar.json pointed at the old function. No behaviour change on the model path: transcripts byte-identical to the previous commit at 60/180 s (granite_nar) and 120 s (qwen3_asr), ctest 89/89. --- src/arch/granite_nar/encoder.cpp | 8 +- src/arch/granite_nar/encoder.h | 5 +- src/arch/granite_nar/model.cpp | 7 +- src/arch/qwen3_asr/encoder.cpp | 24 +-- src/arch/qwen3_asr/encoder.h | 14 +- src/arch/qwen3_asr/model.cpp | 3 +- tests/CMakeLists.txt | 17 +++ tests/granite_nar_bpe_pool_unit.cpp | 219 ++++++++++++++++++++++++++++ tests/tolerances/granite_nar.json | 2 +- 9 files changed, 265 insertions(+), 34 deletions(-) create mode 100644 tests/granite_nar_bpe_pool_unit.cpp diff --git a/src/arch/granite_nar/encoder.cpp b/src/arch/granite_nar/encoder.cpp index 4fc3c456..ea7025f5 100644 --- a/src/arch/granite_nar/encoder.cpp +++ b/src/arch/granite_nar/encoder.cpp @@ -527,12 +527,15 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // Collapse for the in-graph pooled BPE head. -void collapse_bpe_window_ids(const std::vector & window_ids, +bool collapse_bpe_window_ids(const std::vector & window_ids, const std::vector & window_mass, int blank_id, std::vector & out_token_ids) { out_token_ids.clear(); - const size_t n_windows = std::min(window_ids.size(), window_mass.size()); + if (window_ids.size() != window_mass.size()) { + return false; + } + const size_t n_windows = window_ids.size(); out_token_ids.reserve(n_windows); // Same greedy + collapse-repeats + drop-blanks as // compute_bpe_ctc_initial_hypothesis, including the all-blank-window rule @@ -546,6 +549,7 @@ void collapse_bpe_window_ids(const std::vector & window_ids, } prev = id; } + return true; } // Host-side BPE CTC pool + greedy decode (reference formulation). diff --git a/src/arch/granite_nar/encoder.h b/src/arch/granite_nar/encoder.h index ec2ba6fa..9cdc88a9 100644 --- a/src/arch/granite_nar/encoder.h +++ b/src/arch/granite_nar/encoder.h @@ -108,7 +108,10 @@ constexpr int kBpeTileWindows = 256; // Greedy collapse of the in-graph pooled BPE head: one argmax id per pool // window and the window's total importance. Windows whose importance sum is // <= 1e-9 emit blank, exactly as compute_bpe_ctc_initial_hypothesis does. -void collapse_bpe_window_ids(const std::vector & window_ids, +// Returns false (and leaves out_token_ids empty) if the two vectors disagree +// in length, which can only be a graph/caller bug: the tiles cover every +// window and both are read from tensors sized eb.n_bpe_windows. +bool collapse_bpe_window_ids(const std::vector & window_ids, const std::vector & window_mass, int blank_id, std::vector & out_token_ids); diff --git a/src/arch/granite_nar/model.cpp b/src/arch/granite_nar/model.cpp index 63177f71..15920eb3 100644 --- a/src/arch/granite_nar/model.cpp +++ b/src/arch/granite_nar/model.cpp @@ -537,7 +537,12 @@ transcribe_status run(transcribe_session * ctx_base, window_ids.resize(off + static_cast(ids->ne[0])); ggml_backend_tensor_get(ids, window_ids.data() + off, 0, static_cast(ids->ne[0]) * sizeof(int32_t)); } - collapse_bpe_window_ids(window_ids, window_mass, /*blank_id=*/cm->hparams.enc_bpe_blank_id, hyp_ids); + if (!collapse_bpe_window_ids(window_ids, window_mass, /*blank_id=*/cm->hparams.enc_bpe_blank_id, hyp_ids)) { + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "granite_nar run: BPE head emitted %zu window ids but %zu window masses", + window_ids.size(), window_mass.size()); + return TRANSCRIBE_ERR_BACKEND; + } } if (hyp_ids.empty()) { // No tokens — emit empty transcript and return. diff --git a/src/arch/qwen3_asr/encoder.cpp b/src/arch/qwen3_asr/encoder.cpp index 8bb9c6a2..ec339eda 100644 --- a/src/arch/qwen3_asr/encoder.cpp +++ b/src/arch/qwen3_asr/encoder.cpp @@ -71,28 +71,6 @@ std::vector build_sinusoid_pe(int32_t d_model, int32_t length, double max return pe; } -// The single-shot graph no longer takes this mask: it is all zeros, which -// is exactly soft_max_ext / flash_attn_ext with a null mask, and as a dense -// [T_enc, T_enc] input it was a T^2 host + device buffer (~1 GB each at -// 20 minutes). Kept as the documented statement of the attention pattern. -std::vector build_cu_seqlens_mask(const EncoderTiming & t, const QwenAsrHParams & hp) { - (void) hp; - const int32_t T = t.T_enc; - // Full attention over the valid aftercnn rows. The measured - // upstream reference (qwen_asr 0.0.6 + transformers eager / sdpa) - // ignores cu_seqlens and runs unmasked bidirectional attention over - // the post-pad-select tensor. vLLM's flash-attn-2 path honors - // cu_seqlens and chunks at window_aftercnn, but its LibriSpeech WER - // is worse than the eager path in our measurements; we follow the - // eager reference so that transcribe.cpp's greedy decode matches - // the per-tensor dumps byte-for-byte after the pad-row trim in - // build_enc_graph. - // - // Zero-filled mask = softmax(scale * QK) with no bias, which is - // identical to soft_max_ext(nullptr) and matches eager semantics. - return std::vector(static_cast(T) * T, 0.0f); -} - // Graph construction namespace { @@ -347,7 +325,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, named(eb.pos_emb_in, "enc.pos_emb.in"); ggml_set_input(eb.pos_emb_in); - // No attention mask: full unmasked attention (see build_cu_seqlens_mask). + // No attention mask: full unmasked attention (see encoder.h). eb.mask_in = nullptr; // ----- Subsample: 3x Conv2d + GELU + conv_out linear, per chunk group ----- diff --git a/src/arch/qwen3_asr/encoder.h b/src/arch/qwen3_asr/encoder.h index 92db6102..2e856acd 100644 --- a/src/arch/qwen3_asr/encoder.h +++ b/src/arch/qwen3_asr/encoder.h @@ -6,7 +6,7 @@ // .Qwen3ASRAudioEncoder. // // Attention: we match the reference's eager full-bidirectional attention over -// the pad-trimmed sequence (see build_cu_seqlens_mask), not vLLM's +// the pad-trimmed sequence (full, unmasked attention), not vLLM's // cu_seqlens chunking. // // Shape conventions (ggml fast-to-slow ne[]): @@ -23,8 +23,13 @@ // The last chunk's aftercnn trailing pad rows are dropped // in the graph (matches reference's // `padded_embed[padded_mask_after_cnn]` selection). -// attn mask : [T_enc, T_enc] additive. All zeros — full -// bidirectional attention over the valid aftercnn rows. +// attn mask : none. The single-shot encoder attends unmasked over the +// whole (pad-trimmed) sequence, which is what an all-zero +// [T_enc, T_enc] additive mask means, so the graph passes +// mask=nullptr rather than materialising a T^2 buffer. +// (The reference's measured eager path ignores cu_seqlens and +// runs plain bidirectional attention over the valid aftercnn +// rows; vLLM's flash-attn-2 windowing scores worse WER.) // output : [output_dim, T_enc] #pragma once @@ -67,7 +72,6 @@ std::vector build_sinusoid_pe(int32_t d_model, int32_t length, double max // Additive attention bias [T_enc, T_enc]. All zeros — full bidirectional // attention over the pad-trimmed sequence (the eager baseline, not vLLM's // cu_seqlens block-diagonal pattern). -std::vector build_cu_seqlens_mask(const EncoderTiming & t, const QwenAsrHParams & hp); struct EncoderDumps { ggml_tensor * mel_in = nullptr; // graph input @@ -82,7 +86,7 @@ struct EncoderDumps { struct EncoderBuild { ggml_tensor * mel_in = nullptr; // [mel_per_chunk, n_mels, 1, n_chunks] ggml_tensor * pos_emb_in = nullptr; // [d_model, per_chunk_aftercnn] - ggml_tensor * mask_in = nullptr; // [T_enc, T_enc] + ggml_tensor * mask_in = nullptr; // always null: full attention ggml_tensor * out = nullptr; // [output_dim, T_enc] EncoderDumps dumps{}; ggml_cgraph * graph = nullptr; diff --git a/src/arch/qwen3_asr/model.cpp b/src/arch/qwen3_asr/model.cpp index 81ca130e..145136b6 100644 --- a/src/arch/qwen3_asr/model.cpp +++ b/src/arch/qwen3_asr/model.cpp @@ -699,7 +699,8 @@ transcribe_status run(transcribe_session * session, } // No attention mask: the encoder attends unmasked over the whole - // utterance (see build_cu_seqlens_mask), so there is nothing to upload. + // utterance (an all-zero additive mask is the same thing), so there is + // nothing to upload. transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 02286afa..110645b0 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -265,6 +265,23 @@ transcribe_apply_warnings(transcribe_prefill_chunk_mask_unit) add_test(NAME transcribe_prefill_chunk_mask_unit COMMAND transcribe_prefill_chunk_mask_unit) +# ----------------------------------------------------------------------------- +# granite_nar BPE CTC pool: in-graph pooled head vs the frame-level reference +# (pure host, no model) + +add_executable(transcribe_granite_nar_bpe_pool_unit + granite_nar_bpe_pool_unit.cpp) + +target_link_libraries(transcribe_granite_nar_bpe_pool_unit PRIVATE transcribe ggml) + +target_include_directories(transcribe_granite_nar_bpe_pool_unit PRIVATE + ${CMAKE_SOURCE_DIR}/src) + +transcribe_apply_warnings(transcribe_granite_nar_bpe_pool_unit) + +add_test(NAME transcribe_granite_nar_bpe_pool_unit + COMMAND transcribe_granite_nar_bpe_pool_unit) + # ----------------------------------------------------------------------------- # Per-run decode budget rule (pure host, no model) diff --git a/tests/granite_nar_bpe_pool_unit.cpp b/tests/granite_nar_bpe_pool_unit.cpp new file mode 100644 index 00000000..7c65024e --- /dev/null +++ b/tests/granite_nar_bpe_pool_unit.cpp @@ -0,0 +1,219 @@ +// granite_nar BPE CTC pool: the in-graph pooled head against the frame-level +// reference formulation. Pure host, no model. +// +// run() does not decode frame-level BPE logits any more (that tensor was +// [100352, T_enc] f32, 6.1 GB at 5 minutes of audio, then copied to the host +// whole). The head is affine and the pool weights sum to 1, so +// sum_t w_t (W h_t + b) = W (sum_t w_t h_t) + b +// and the graph pools the 1024-wide hidden states first, runs the head on one +// column per window, and takes the argmax there. This test is what keeps that +// restructure honest: it builds the same head both ways and requires the two +// decodes to agree token for token, including the all-blank-window rule and +// both blank-id schemes. +// +// It also pins the failure mode the caller relies on: mismatched id/mass +// lengths are an error, not a silently shortened hypothesis. + +#include "arch/granite_nar/encoder.h" + +#include +#include +#include +#include + +namespace { + +int g_failures = 0; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond); \ + ++g_failures; \ + } \ + } while (0) + +// Deterministic, well-spread values without pulling in a RNG. +float pseudo(int i) { + const float s = std::sin(static_cast(i) * 12.9898f + 78.233f) * 43758.5453f; + return s - std::floor(s) - 0.5f; +} + +struct Head { + int dim; + int vocab; + std::vector w; // [vocab, dim] + std::vector b; // [vocab] + + Head(int dim_, int vocab_) : dim(dim_), vocab(vocab_), w(static_cast(dim_) * vocab_), b(vocab_) { + for (size_t i = 0; i < w.size(); ++i) { + w[i] = pseudo(static_cast(i)); + } + for (int v = 0; v < vocab; ++v) { + b[static_cast(v)] = pseudo(100000 + v) * 0.1f; + } + } + + void apply(const std::vector & x, int n_cols, std::vector & out) const { + // out[col][v] = sum_i w[v][i] * x[i][col] + b[v] + out.assign(static_cast(n_cols) * vocab, 0.0f); + for (int v = 0; v < vocab; ++v) { + for (int c = 0; c < n_cols; ++c) { + float acc = b[static_cast(v)]; + for (int i = 0; i < dim; ++i) { + acc += w[static_cast(v) * dim + i] * x[static_cast(c) * dim + i]; + } + out[static_cast(c) * vocab + v] = acc; + } + } + } +}; + +// One scenario: `hidden` is [T, dim] row-major (frame-major), importance is +// [T]. Runs the reference (frame-level logits, then pool + collapse) and the +// graph's shape (pool hidden, then head + argmax + collapse) and compares. +void check_equivalence(const char * what, + const Head & head, + const std::vector & hidden, + const std::vector & importance, + int pool_window, + int blank_id) { + const int T = static_cast(importance.size()); + const int n_win = (T + pool_window - 1) / pool_window; + + // --- Reference: frame-level logits, host pool + greedy (encoder.cpp) --- + std::vector frame_logits; + head.apply(hidden, T, frame_logits); + std::vector ref_ids; + transcribe::granite_nar::compute_bpe_ctc_initial_hypothesis(importance, frame_logits, head.vocab, T, pool_window, + blank_id, ref_ids); + + // --- Graph shape: importance-weighted hidden pool, then head + argmax --- + std::vector pooled(static_cast(n_win) * head.dim, 0.0f); + std::vector mass(static_cast(n_win), 0.0f); + std::vector window_ids(static_cast(n_win), blank_id); + for (int w = 0; w < n_win; ++w) { + const int t0 = w * pool_window; + const int t1 = std::min(t0 + pool_window, T); + for (int t = t0; t < t1; ++t) { + mass[static_cast(w)] += importance[static_cast(t)]; + } + // Softmax mid-blank CT counts are in (0, 1], so a window mass is + // never 0 unless every frame in it is blank; the graph divides by a + // clamped mass and the host then treats the window as blank. + const float denom = mass[static_cast(w)] > 1e-9f ? mass[static_cast(w)] : 1e-9f; + for (int t = t0; t < t1; ++t) { + const float wt = importance[static_cast(t)] / denom; + for (int i = 0; i < head.dim; ++i) { + pooled[static_cast(w) * head.dim + i] += + wt * hidden[static_cast(t) * head.dim + i]; + } + } + } + std::vector win_logits; + head.apply(pooled, n_win, win_logits); + for (int w = 0; w < n_win; ++w) { + int best = 0; + float best_v = win_logits[static_cast(w) * head.vocab]; + for (int v = 1; v < head.vocab; ++v) { + const float cand = win_logits[static_cast(w) * head.vocab + v]; + if (cand > best_v) { + best_v = cand; + best = v; + } + } + window_ids[static_cast(w)] = best; + } + std::vector got_ids; + CHECK(transcribe::granite_nar::collapse_bpe_window_ids(window_ids, mass, blank_id, got_ids)); + + if (ref_ids != got_ids) { + ++g_failures; + std::fprintf(stderr, "FAIL %s: %zu vs %zu tokens differ (blank_id=%d, T=%d, pool=%d)\n", what, ref_ids.size(), + got_ids.size(), blank_id, T, pool_window); + const size_t n = std::min(ref_ids.size(), got_ids.size()); + for (size_t i = 0; i < n; ++i) { + if (ref_ids[i] != got_ids[i]) { + std::fprintf(stderr, " first diff at %zu: reference %d, pooled %d\n", i, ref_ids[i], got_ids[i]); + break; + } + } + // The pooled head is 4x cheaper but does not reproduce per-window + // division in the same summation order, so allow a tiny numerical + // disagreement to be visible rather than silently tolerated: the + // test's inputs are well-separated integers, so exact equality is the + // right bar. Print the reference to make a real drift obvious. + std::fprintf(stderr, " reference:"); + for (int32_t id : ref_ids) { + std::fprintf(stderr, " %d", id); + } + std::fprintf(stderr, "\n"); + } +} + +} // namespace + +int main() { + // A tiny head so the arithmetic stays exact in f32: values are integers, + // a window's weighted mean is representable, and argmax has no ties. + Head head(/*dim=*/4, /*vocab=*/16); + + // T = 10, P = 4 -> 3 windows, the last one partial (2 frames). + const int T = 10; + const int P = 4; + std::vector hidden(static_cast(T) * head.dim); + for (int t = 0; t < T; ++t) { + for (int i = 0; i < head.dim; ++i) { + hidden[static_cast(t) * head.dim + i] = static_cast(1 + ((t * 7 + i * 3) % 11)); + } + } + + // All frames voiced, weights 1/2/4/8 cycling: exercises the weighting. + std::vector voiced(static_cast(T)); + for (int t = 0; t < T; ++t) { + voiced[static_cast(t)] = static_cast(1 << (t % 4)); + } + + // Scheme A (blank_id == 0): synthetic blank at channel 0, ids shifted. + check_equivalence("voiced/blank=0", head, hidden, voiced, P, /*blank_id=*/0); + // Scheme B (blank_id == 100257, BOS): channels are LLM ids directly. + check_equivalence("voiced/blank=BOS", head, hidden, voiced, P, /*blank_id=*/100257); + + // A silent window in the middle (mass 0) plus a partial tail: the + // reference emits blank there and the collapse must do the same. + std::vector with_gap = voiced; + for (int t = P; t < 2 * P; ++t) { + with_gap[static_cast(t)] = 0.0f; + } + check_equivalence("gap/blank=0", head, hidden, with_gap, P, /*blank_id=*/0); + check_equivalence("gap/blank=BOS", head, hidden, with_gap, P, /*blank_id=*/100257); + + // Collapse rules: adjacent duplicate windows collapse, blanks are dropped. + { + const std::vector ids = {3, 3, 0, 7, 0, 0, 7}; + const std::vector mass = {1, 1, 1, 1, 1, 1, 1}; + std::vector out; + CHECK(transcribe::granite_nar::collapse_bpe_window_ids(ids, mass, /*blank_id=*/0, out)); + // Collapse runs on the window ids (before the shift): 3,3 -> one 3; + // 0 is blank; 7 survives both times because a blank resets "prev". + // Scheme A then subtracts 1 from each emitted id. + const std::vector expect = {2, 6, 6}; + CHECK(out == expect); + } + + // A length mismatch is a caller bug: refuse instead of truncating. + { + const std::vector ids = {1, 2, 3}; + const std::vector mass = {1, 1}; + std::vector out = {42}; + CHECK(!transcribe::granite_nar::collapse_bpe_window_ids(ids, mass, /*blank_id=*/0, out)); + CHECK(out.empty()); + } + + if (g_failures == 0) { + std::fprintf(stdout, "granite_nar_bpe_pool_unit: ok\n"); + return 0; + } + std::fprintf(stderr, "granite_nar_bpe_pool_unit: %d failures\n", g_failures); + return 1; +} diff --git a/tests/tolerances/granite_nar.json b/tests/tolerances/granite_nar.json index d87e3cc2..3bd8106e 100644 --- a/tests/tolerances/granite_nar.json +++ b/tests/tolerances/granite_nar.json @@ -54,7 +54,7 @@ "argmax index directly. The older 7d20732d snapshot used bpe_output_dim", "= vocab_size + 1 with blank at channel 0 and `argmax - 1` shift. The", "converter writes the actual blank_token_id as a new KV so C++ honors", - "the right scheme — see encoder.cpp::compute_bpe_ctc_initial_hypothesis.", + "the right scheme — see collapse_bpe_window_ids in encoder.cpp.", "", "All numbers measured with `validate.py all --family granite_nar --variant", "granite-speech-4.1-2b-nar`. Recipe: max(1.5x observed, 1e-6)." From f1b44a5d32d4fa08848c0ba5f0b51f0fc0486856 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Wed, 21 Oct 2026 16:40:27 +0800 Subject: [PATCH 5/6] qwen3_asr: run the conv frontend in fixed-size passes so long audio cannot overflow the graph --- docs/environment-variables.md | 1 + src/arch/qwen3_asr/encoder.cpp | 123 ++++++++++---------- src/arch/qwen3_asr/encoder.h | 69 ++++++++--- src/arch/qwen3_asr/model.cpp | 202 ++++++++++++++++++++++++++------- 4 files changed, 279 insertions(+), 116 deletions(-) diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 9fa946e6..04c2a233 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -36,6 +36,7 @@ of tests. | `TRANSCRIBE_PERF_DEBUG` | Print a per-stage timing breakdown to stderr (DEBUG log) on the families that profile (`cohere`, `granite`, `canary`, `canary_qwen`, `moonshine`, `moonshine_streaming`, `moss`, `qwen3_asr`, `whisper`). For whisper, a value containing `cpu` or `all` additionally prints the CPU sub-section breakdown. | | `TRANSCRIBE_VOXTRAL_REALTIME_STREAM_TIMING` | Print a per-component streaming wall-time breakdown at stream finalize (voxtral_realtime). | | `TRANSCRIBE_MOONSHINE_STREAMING_ENC_WINDOW=` | Encoder frames emitted per window by the moonshine_streaming one-shot / batch encoder (default 1024, ~20 s). Windowing bounds encoder memory; it is mathematically exact (each window carries the full sliding-window context) and float-level differences stay at the backend's reduction-order noise. Set small to exercise the windowed path on short clips. Read at each run. | +| `TRANSCRIBE_QWEN3_ASR_CONV_PASS=` | Conv-subsampler chunks (1 s of audio each) per subsample-graph run in the qwen3_asr encoder (default 512, ~8.5 min). Longer inputs run the subsampler in passes so its graph stays a fixed size; chunks are independent, so the value does not change the result (bit-identical on CPU; GPU backends may differ at reduction-order noise for values that are not a multiple of 16). Set small to exercise the multi-pass path on short clips. Read at each run. | | `TRANSCRIBE_TEST_DEV_INIT_THROW=` | Fault injection: backend device init (`ggml_backend_dev_init`) throws for devices whose name contains `` (`*` matches every device). Exercises throw → skip → CPU-fallback in backend probing; an explicit backend request fails with `TRANSCRIBE_ERR_BACKEND`. Used by `backend_init_throw_unit` and `scripts/ci/vulkan_degradation_check.py`. | | `TRANSCRIBE_TEST_TEARDOWN_THROW` | Fault injection: any non-empty value injects a throw after each real free inside the `transcribe::safe_*` teardown wrappers, proving containment without leaking the handle. Used by `teardown_safety_unit`. | diff --git a/src/arch/qwen3_asr/encoder.cpp b/src/arch/qwen3_asr/encoder.cpp index ec339eda..f1ecc006 100644 --- a/src/arch/qwen3_asr/encoder.cpp +++ b/src/arch/qwen3_asr/encoder.cpp @@ -108,18 +108,6 @@ ggml_tensor * linear(ggml_context * ctx, ggml_tensor * x, ggml_tensor * w, ggml_ return y; } -// Conv chunks per subsample group. The 3x Conv2d frontend is independent per -// n_window*2-mel chunk, but running every chunk in one batch makes its -// activations linear in audio length with a huge constant: conv0's output is -// [50, 64, 480] f32 (6.1 MB) per chunk and conv1's im2col another 6.9 MB, so -// ~1.2 GB of scheduler scratch per minute of audio. Grouping bounds that to a -// fixed size: each group's intermediates die before the next group's are -// allocated, and only the small [d_model, T_per_chunk, g] outputs are -// concatenated. The reference does the same thing with conv_chunksize (500); -// ours is smaller because the scheduler commits the whole group at once. -// Per-chunk results do not depend on the group size. -constexpr int64_t kConvChunkGroup = 16; - // Conv subsample + conv_out (no PE) for `mel` with ne=[mel_per_chunk, // n_mels, 1, n] (any n). Returns [d_model, T_per_chunk, n], or nullptr if the // conv output width does not match T_per_chunk. @@ -192,7 +180,7 @@ ggml_tensor * subsample_chunks(ggml_context * ctx, const int64_t n_chunks = mel->ne[3]; ggml_tensor * out = nullptr; for (int64_t c0 = 0; c0 < n_chunks; c0 += kConvChunkGroup) { - const int64_t n = std::min(kConvChunkGroup, n_chunks - c0); + const int64_t n = std::min(kConvChunkGroup, n_chunks - c0); ggml_tensor * g = (n == n_chunks) ? mel : ggml_view_4d(ctx, mel, mel->ne[0], mel->ne[1], mel->ne[2], n, mel->nb[1], mel->nb[2], mel->nb[3], static_cast(c0) * mel->nb[3]); @@ -272,6 +260,42 @@ ggml_tensor * build_enc_block(ggml_context * ctx, } // namespace +SubsampleBuild build_subsample_graph(ggml_context * ctx, + const QwenAsrWeights & weights, + const QwenAsrHParams & hp, + int n_chunks) { + SubsampleBuild sb{}; + if (ctx == nullptr || n_chunks <= 0) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "qwen3_asr subsample: invalid arg (ctx=%p, n_chunks=%d)", + static_cast(ctx), n_chunks); + return sb; + } + + const int64_t mel_per_chunk = hp.enc_n_window * 2; + const int64_t T_per_chunk = aftercnn_len(static_cast(mel_per_chunk)); + + sb.mel_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, mel_per_chunk, hp.enc_num_mel_bins, 1, n_chunks); + named(sb.mel_in, "enc.mel_in"); + ggml_set_input(sb.mel_in); + + sb.out = subsample_chunks(ctx, weights, hp.enc_downsample_hidden, T_per_chunk, sb.mel_in); + if (sb.out == nullptr) { + return sb; + } + named(sb.out, "enc.subsample.pass"); + ggml_set_output(sb.out); + + // ~36 nodes per group (3 convs + GELUs + projection + concat). + const int n_groups = (n_chunks + kConvChunkGroup - 1) / kConvChunkGroup; + sb.graph = ggml_new_graph_custom(ctx, /*size=*/64 + 48 * n_groups, /*grads=*/false); + if (sb.graph == nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "qwen3_asr subsample: ggml_new_graph_custom failed"); + return sb; + } + ggml_build_forward_expand(sb.graph, sb.out); + return sb; +} + EncoderBuild build_encoder_graph(ggml_context * ctx, const QwenAsrWeights & weights, const QwenAsrHParams & hp, @@ -300,26 +324,21 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // shape see T_enc (not T_enc_padded). This matches the reference // byte-for-byte at the pos_add.out boundary and upward. - const int64_t d_model = hp.enc_d_model; - const int64_t n_heads = hp.enc_n_heads; - const int64_t ds_h = hp.enc_downsample_hidden; - const int64_t n_mels = hp.enc_num_mel_bins; - const int64_t mel_per_chunk = timing.mel_per_chunk; - const int64_t n_chunks = timing.n_chunks; - const int64_t T_per_chunk = timing.per_chunk_aftercnn; - const int64_t T_enc_padded = timing.T_enc_padded; - const int64_t T_enc = timing.T_enc; + const int64_t d_model = hp.enc_d_model; + const int64_t n_heads = hp.enc_n_heads; + const int64_t n_chunks = timing.n_chunks; + const int64_t T_per_chunk = timing.per_chunk_aftercnn; + const int64_t T_enc_padded = timing.T_enc_padded; + const int64_t T_enc = timing.T_enc; // ----- Graph inputs ----- - // mel layout: ggml fast-to-slow ne=[mel_per_chunk, n_mels, 1, n_chunks] - // ≡ [W=time, H=n_mels, C=1, N=n_chunks]. ggml_conv_2d with a - // symmetric 3×3 kernel and symmetric stride/pad is invariant under - // a W/H swap, so this is equivalent to Cohere's time-along-H - // convention. - eb.mel_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, mel_per_chunk, n_mels, 1, n_chunks); - named(eb.mel_in, "enc.mel_in"); - ggml_set_input(eb.mel_in); - eb.dumps.mel_in = eb.mel_in; + // Conv subsample output [d_model, T_per_chunk, n_chunks], computed by + // build_subsample_graph (pre-PE). + eb.subsample_in = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_model, T_per_chunk, n_chunks); + named(eb.subsample_in, "enc.subsample.out"); + ggml_set_input(eb.subsample_in); + eb.dumps.subsample_out = eb.subsample_in; + transcribe::debug::mark_tensor_for_dump(eb.subsample_in); eb.pos_emb_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_model, T_per_chunk); named(eb.pos_emb_in, "enc.pos_emb.in"); @@ -328,19 +347,9 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // No attention mask: full unmasked attention (see encoder.h). eb.mask_in = nullptr; - // ----- Subsample: 3x Conv2d + GELU + conv_out linear, per chunk group ----- - ggml_tensor * x = subsample_chunks(ctx, weights, ds_h, T_per_chunk, eb.mel_in); - if (x == nullptr) { - return eb; - } - // ne = [d_model, T_per_chunk, n_chunks] - named(x, "enc.subsample.out"); - eb.dumps.subsample_out = x; - transcribe::debug::mark_tensor_for_dump(x); - // Add sinusoidal PE (broadcast across n_chunks). pos_emb_in has // ne=[d_model, T_per_chunk]; ggml_add broadcasts missing axes. - x = ggml_add(ctx, x, eb.pos_emb_in); + ggml_tensor * x = ggml_add(ctx, eb.subsample_in, eb.pos_emb_in); // Flatten batch: [d_model, T_per_chunk, n_chunks] -> [d_model, T_enc_padded]. // Row-major semantics: position r in the flattened sequence is @@ -458,15 +467,12 @@ EncoderBuildBatched build_encoder_graph_batched(ggml_context * ctx, return eb; } - const int64_t d_model = hp.enc_d_model; - const int64_t n_heads = hp.enc_n_heads; - const int64_t ds_h = hp.enc_downsample_hidden; - const int64_t n_mels = hp.enc_num_mel_bins; - const int64_t mel_per_chunk = hp.enc_n_window * 2; - const int64_t T_per_chunk = aftercnn_len(static_cast(mel_per_chunk)); - const int64_t B = n_batch; - const int64_t N = B * n_chunks_max; // packed conv batch - const int64_t T_pad_max = n_chunks_max * T_per_chunk; + const int64_t d_model = hp.enc_d_model; + const int64_t n_heads = hp.enc_n_heads; + const int64_t T_per_chunk = aftercnn_len(hp.enc_n_window * 2); + const int64_t B = n_batch; + const int64_t N = B * n_chunks_max; // packed chunk axis + const int64_t T_pad_max = n_chunks_max * T_per_chunk; eb.n_batch = n_batch; eb.n_chunks_max = n_chunks_max; @@ -474,9 +480,10 @@ EncoderBuildBatched build_encoder_graph_batched(ggml_context * ctx, eb.T_pad_max = static_cast(T_pad_max); // ----- Graph inputs ----- - eb.mel_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, mel_per_chunk, n_mels, 1, N); - named(eb.mel_in, "enc.mel_in"); - ggml_set_input(eb.mel_in); + // Conv subsample output for all N packed chunks (build_subsample_graph). + eb.subsample_in = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_model, T_per_chunk, N); + named(eb.subsample_in, "enc.subsample.out"); + ggml_set_input(eb.subsample_in); eb.pos_emb_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_model, T_per_chunk); named(eb.pos_emb_in, "enc.pos_emb.in"); @@ -487,14 +494,8 @@ EncoderBuildBatched build_encoder_graph_batched(ggml_context * ctx, named(eb.mask_in, "enc.attn_mask.in"); ggml_set_input(eb.mask_in); - // ----- Subsample: 3x Conv2d + GELU + conv_out (per-chunk over N = B*n_chunks_max) ----- - ggml_tensor * x = subsample_chunks(ctx, weights, ds_h, T_per_chunk, eb.mel_in); // [d_model, T_per_chunk, N] - if (x == nullptr) { - return eb; - } - // Add sinusoidal PE (broadcast across the N = B*n_chunks_max axis). - x = ggml_add(ctx, x, eb.pos_emb_in); + ggml_tensor * x = ggml_add(ctx, eb.subsample_in, eb.pos_emb_in); // Reorganise N=(b*n_chunks_max + c) -> [d_model, T_per_chunk, n_chunks_max, B], // then flatten (T_per_chunk, n_chunks_max) chunk-major into the sequence: diff --git a/src/arch/qwen3_asr/encoder.h b/src/arch/qwen3_asr/encoder.h index 2e856acd..fbaec628 100644 --- a/src/arch/qwen3_asr/encoder.h +++ b/src/arch/qwen3_asr/encoder.h @@ -9,15 +9,21 @@ // the pad-trimmed sequence (full, unmasked attention), not vLLM's // cu_seqlens chunking. // +// The conv subsampler and the transformer are two graphs: the subsampler +// (build_subsample_graph) runs over up to kConvPassChunks chunks at a time, +// and its outputs are concatenated host-side into the encoder graph's input. +// // Shape conventions (ggml fast-to-slow ne[]): // -// mel_in : [mel_per_chunk, n_mels, 1, n_chunks] -// per-chunk batched input. Chunks with fewer real mel -// frames than mel_per_chunk are zero-padded by the caller. -// after 3x conv: [F_ds=n_mels/8, T_ds=mel_per_chunk/8, downsample_hidden, n_chunks] +// mel_in : [mel_per_chunk, n_mels, 1, n] +// subsample-graph input, n <= kConvPassChunks chunks. +// Chunks with fewer real mel frames than mel_per_chunk +// are zero-padded by the caller. +// after 3x conv: [F_ds=n_mels/8, T_ds=mel_per_chunk/8, downsample_hidden, n] // (ggml conv_2d convention swaps PyTorch's H/W, which is // a no-op for the square kernel + symmetric stride.) -// after conv_out + PE: [d_model, per_chunk_aftercnn, n_chunks] +// subsample_in: [d_model, per_chunk_aftercnn, n_chunks] +// conv_out output, the encoder graph's input; PE added there. // after pad-select: [d_model, T_enc] where T_enc = // (n_chunks-1) * per_chunk_aftercnn + last_chunk_aftercnn. // The last chunk's aftercnn trailing pad rows are dropped @@ -73,9 +79,40 @@ std::vector build_sinusoid_pe(int32_t d_model, int32_t length, double max // attention over the pad-trimmed sequence (the eager baseline, not vLLM's // cu_seqlens block-diagonal pattern). +// Conv chunks per group inside a subsample graph. The 3x Conv2d frontend is +// independent per n_window*2-mel chunk, but its activations are large: +// conv0's output is [50, 64, 480] f32 (6.1 MB) per chunk and conv1's im2col +// another 6.9 MB, ~1.2 GB of scheduler scratch per minute of audio if every +// chunk runs at once. The graph runs the chunks in groups chained by concat, +// so each group's intermediates retire before the next group's are +// allocated and scratch stays at one group's worth. The reference does the +// same with conv_chunksize (500); ours is smaller because the scheduler +// commits the whole group at once. Per-chunk results do not depend on the +// group size. +constexpr int kConvChunkGroup = 16; + +// Conv chunks per subsample graph run (~8.5 minutes of audio). Each group +// adds ~36 graph nodes, so one graph over the whole input would overflow +// its node budget on long audio; longer inputs re-run a graph of this many +// chunks per pass instead. Large enough that typical inputs take one pass, +// since each pass costs a backend sync. +constexpr int kConvPassChunks = 512; + +struct SubsampleBuild { + ggml_tensor * mel_in = nullptr; // [mel_per_chunk, n_mels, 1, n] + ggml_tensor * out = nullptr; // [d_model, per_chunk_aftercnn, n] + ggml_cgraph * graph = nullptr; +}; + +// Conv subsample + conv_out linear (no PE) over n chunks, in groups of +// kConvChunkGroup. +SubsampleBuild build_subsample_graph(ggml_context * ctx, + const QwenAsrWeights & weights, + const QwenAsrHParams & hp, + int n_chunks); + struct EncoderDumps { - ggml_tensor * mel_in = nullptr; // graph input - ggml_tensor * subsample_out = nullptr; // post conv_out linear, pre-PE + ggml_tensor * subsample_out = nullptr; // post conv_out linear, pre-PE (graph input) ggml_tensor * pos_add_out = nullptr; // post-PE, flattened and pad-trimmed to [d_model, T_enc] ggml_tensor * block_0_out = nullptr; ggml_tensor * block_last_out = nullptr; @@ -84,9 +121,9 @@ struct EncoderDumps { }; struct EncoderBuild { - ggml_tensor * mel_in = nullptr; // [mel_per_chunk, n_mels, 1, n_chunks] - ggml_tensor * pos_emb_in = nullptr; // [d_model, per_chunk_aftercnn] - ggml_tensor * mask_in = nullptr; // always null: full attention + ggml_tensor * subsample_in = nullptr; // [d_model, per_chunk_aftercnn, n_chunks] + ggml_tensor * pos_emb_in = nullptr; // [d_model, per_chunk_aftercnn] + ggml_tensor * mask_in = nullptr; // unused (null): full unmasked attention ggml_tensor * out = nullptr; // [output_dim, T_enc] EncoderDumps dumps{}; ggml_cgraph * graph = nullptr; @@ -104,11 +141,11 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // --------------------------------------------------------------------------- struct EncoderBuildBatched { - ggml_tensor * mel_in = nullptr; // [mel_per_chunk, n_mels, 1, B*n_chunks_max] - ggml_tensor * pos_emb_in = nullptr; // [d_model, per_chunk_aftercnn] - ggml_tensor * mask_in = nullptr; // [T_pad_max, T_pad_max, 1, B] key-pad - ggml_tensor * out = nullptr; // [output_dim, T_pad_max, B] - ggml_cgraph * graph = nullptr; + ggml_tensor * subsample_in = nullptr; // [d_model, per_chunk_aftercnn, B*n_chunks_max] + ggml_tensor * pos_emb_in = nullptr; // [d_model, per_chunk_aftercnn] + ggml_tensor * mask_in = nullptr; // [T_pad_max, T_pad_max, 1, B] key-pad + ggml_tensor * out = nullptr; // [output_dim, T_pad_max, B] + ggml_cgraph * graph = nullptr; int n_batch = 0; int n_chunks_max = 0; @@ -119,7 +156,7 @@ struct EncoderBuildBatched { // Build one encoder graph that processes B utterances in parallel on the // batch axis ne[2]. All utterances share the same per_chunk_aftercnn (it // depends only on enc_n_window), so they pack cleanly: utterance b's chunks -// occupy N-indices [b*n_chunks_max, b*n_chunks_max + n_chunks[b]) of mel_in, +// occupy N-indices [b*n_chunks_max, b*n_chunks_max + n_chunks[b]) of subsample_in, // zero-padded to n_chunks_max. The conv subsampler is per-chunk (no cross- // utterance leak); the 18 blocks attend per-utterance under a key-pad mask // (`mask_in[k,q,0,b] = 0 if k < T_enc[b] else -inf`). The real rows of diff --git a/src/arch/qwen3_asr/model.cpp b/src/arch/qwen3_asr/model.cpp index 145136b6..bccb263e 100644 --- a/src/arch/qwen3_asr/model.cpp +++ b/src/arch/qwen3_asr/model.cpp @@ -563,6 +563,75 @@ void pack_mel_chunks(const float * mel, // [n_mels, T_mel] } } +// TRANSCRIBE_QWEN3_ASR_CONV_PASS= overrides kConvPassChunks, so the +// multi-pass path can be exercised (and checked against one pass) on short +// clips. Values below 1 are ignored. +int conv_pass_chunks() { + if (const char * v = transcribe::env::str("TRANSCRIBE_QWEN3_ASR_CONV_PASS"); v != nullptr) { + char * end = nullptr; + const long n = std::strtol(v, &end, 10); + if (end != v && n >= 1 && n <= (1 << 20)) { + return static_cast(n); + } + } + return kConvPassChunks; +} + +// Conv subsampler over `n_chunks` packed mel chunks (pack_mel_chunks layout) +// into out = [d_model, per_chunk_aftercnn, n_chunks], in passes of +// conv_pass_chunks() chunks. Full passes share one graph, allocated once and +// re-run; the remainder gets one graph of its own size. Neither node count +// nor scratch grows with the audio, and chunks are independent, so the +// result matches one graph over every chunk. Built in cc->compute_ctx; +// resets cc->sched, and the caller resets it again before its next graph. +transcribe_status run_subsample(QwenAsrSession * cc, + QwenAsrModel * cm, + const float * mel, + int n_chunks, + std::vector & out) { + const size_t in_chunk = static_cast(cm->hparams.enc_n_window) * 2 * cm->hparams.enc_num_mel_bins; + const size_t out_chunk = static_cast(cm->hparams.enc_d_model) * aftercnn_len(cm->hparams.enc_n_window * 2); + out.resize(out_chunk * n_chunks); + + // Runs chunks [c_begin, c_end) through one graph of `pass` chunks. + auto run_passes = [&](int pass, int c_begin, int c_end) -> transcribe_status { + SubsampleBuild sb = build_subsample_graph(cc->compute_ctx, cm->weights, cm->hparams, pass); + if (sb.graph == nullptr || sb.out == nullptr) { + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_sched_reset(cc->sched); + if (!TRANSCRIBE_ALLOC_SCHED_GRAPH(cc->sched, sb.graph, "qwen3_asr")) { + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "qwen3_asr run: subsample graph allocation failed — out of memory."); + return TRANSCRIBE_ERR_OOM; + } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); + for (int c0 = c_begin; c0 < c_end; c0 += pass) { + ggml_backend_tensor_set(sb.mel_in, mel + static_cast(c0) * in_chunk, 0, ggml_nbytes(sb.mel_in)); + if (const ggml_status gs = ggml_backend_sched_graph_compute(cc->sched, sb.graph); + gs != GGML_STATUS_SUCCESS) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "qwen3_asr run: subsample graph compute failed (%d)", + static_cast(gs)); + return TRANSCRIBE_ERR_BACKEND; + } + ggml_backend_tensor_get(sb.out, out.data() + static_cast(c0) * out_chunk, 0, ggml_nbytes(sb.out)); + } + return TRANSCRIBE_OK; + }; + + const int pass = conv_pass_chunks(); + const int n_full = n_chunks - n_chunks % pass; + if (n_full > 0) { + if (const transcribe_status st = run_passes(pass, 0, n_full); st != TRANSCRIBE_OK) { + return st; + } + } + if (n_full < n_chunks) { + return run_passes(n_chunks - n_full, n_full, n_chunks); + } + return TRANSCRIBE_OK; +} + transcribe_status run(transcribe_session * session, const float * pcm, int n_samples, @@ -635,6 +704,23 @@ transcribe_status run(transcribe_session * session, return TRANSCRIBE_ERR_GGUF; } + // Input-length gate, before encoding: audio + the prompt without system + // context + generation must fit the decoder context window. The system + // context is fitted to what is left (encode_system_context), so this is + // the same bound as the prompt gate below, checked before an over-length + // clip pays for the encoder. + const int ceiling = qwen3_context_ceiling(cc->n_ctx, cm->hparams); + if (const int T_min = prompt_tokens_without_system(timing.T_enc, lang_prefix_ptr); + T_min + k_gen_reserve > ceiling) { + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "qwen3_asr run: input too long — %d audio + %d prompt tokens " + "leave no room for output within the %d-token context (need %d). " + "Shorten the audio (see transcribe_capabilities.max_audio_ms) or " + "split it into segments.", + timing.T_enc, T_min - timing.T_enc, ceiling, T_min + k_gen_reserve); + return TRANSCRIBE_ERR_INPUT_TOO_LONG; + } + // Reset per-call compute state. The phase timers below break out // per-run cost (graph build, sched alloc, uploads, prefill compute) // that the public transcribe_timings (mel/encode/decode) doesn't @@ -665,13 +751,6 @@ transcribe_status run(transcribe_session * session, } } - // Build encoder graph. - EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, cm->hparams, timing, cc->encoder_use_flash); - if (eb.graph == nullptr || eb.out == nullptr) { - return TRANSCRIBE_ERR_GGUF; - } - - // Allocate + compute encoder graph. if (cc->sched == nullptr) { cc->sched = ggml_backend_sched_new(cm->plan.scheduler_list.data(), nullptr, static_cast(cm->plan.scheduler_list.size()), 16384, false, true); @@ -680,6 +759,28 @@ transcribe_status run(transcribe_session * session, return TRANSCRIBE_ERR_BACKEND; } } + + // Conv subsampler, in passes of up to kConvPassChunks chunks. + std::vector subsample_host; + int64_t t_subsample_us = 0; + { + std::vector mel_batched; + pack_mel_chunks(cc->mel_buf.data(), mel_n_mels, mel_n_frames, timing, mel_batched); + const int64_t t_sub0 = ggml_time_us(); + if (const transcribe_status st = run_subsample(cc, cm, mel_batched.data(), timing.n_chunks, subsample_host); + st != TRANSCRIBE_OK) { + return st; + } + t_subsample_us = ggml_time_us() - t_sub0; + } + + // Build encoder graph. + EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, cm->hparams, timing, cc->encoder_use_flash); + if (eb.graph == nullptr || eb.out == nullptr) { + return TRANSCRIBE_ERR_GGUF; + } + + // Allocate + compute encoder graph. ggml_backend_sched_reset(cc->sched); if (!TRANSCRIBE_ALLOC_SCHED_GRAPH(cc->sched, eb.graph, "qwen3_asr")) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -687,10 +788,7 @@ transcribe_status run(transcribe_session * session, return TRANSCRIBE_ERR_OOM; } - // Pack + upload mel. - std::vector mel_batched; - pack_mel_chunks(cc->mel_buf.data(), mel_n_mels, mel_n_frames, timing, mel_batched); - ggml_backend_tensor_set(eb.mel_in, mel_batched.data(), 0, mel_batched.size() * sizeof(float)); + ggml_backend_tensor_set(eb.subsample_in, subsample_host.data(), 0, subsample_host.size() * sizeof(float)); // Positional embedding. { @@ -705,12 +803,12 @@ transcribe_status run(transcribe_session * session, transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); const int64_t t_enc_start = ggml_time_us(); - t_enc_build_us = t_enc_start - t_enc_build_start; + t_enc_build_us = t_enc_start - t_enc_build_start - t_subsample_us; if (const ggml_status gs = ggml_backend_sched_graph_compute(cc->sched, eb.graph); gs != GGML_STATUS_SUCCESS) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "qwen3_asr run: encoder graph compute failed (%d)", static_cast(gs)); return TRANSCRIBE_ERR_BACKEND; } - cc->t_encode_us = ggml_time_us() - t_enc_start; + cc->t_encode_us = ggml_time_us() - t_enc_start + t_subsample_us; // Dump encoder intermediates. auto try_dump = [](const char * name, ggml_tensor * t, const char * stage) { @@ -748,7 +846,6 @@ transcribe_status run(transcribe_session * session, // Prompt construction. std::vector prompt_ids; std::vector audio_positions; - const int ceiling = qwen3_context_ceiling(cc->n_ctx, cm->hparams); std::vector system_ids; if (const transcribe_status st = encode_system_context( cm->tok, params, ceiling - k_gen_reserve - prompt_tokens_without_system(T_enc, lang_prefix_ptr), @@ -1162,13 +1259,17 @@ transcribe_status reset_compute_ctx(QwenAsrSession * cc, int mb) { // Batched encoder: mel (parallel) + one encoder graph over all B utterances // on the batch axis. Fills enc_hosts[b] = [d_enc, T_enc[b]], T_enc[b], valid[b]. -// Real-row outputs are bit-identical to encode_one per utterance. +// Real-row outputs are bit-identical to encode_one per utterance. A row with +// more than T_enc_cap audio tokens is rejected (fail_status INPUT_TOO_LONG) +// before it can size the batch. transcribe_status encode_all_batched(QwenAsrSession * cc, QwenAsrModel * cm, const float * const * pcm, const int * n_samples, int n, + int T_enc_cap, std::vector & valid, + std::vector & fail_status, std::vector> & enc_hosts, std::vector & T_enc_out, int64_t & mel_us, @@ -1211,6 +1312,18 @@ transcribe_status encode_all_batched(QwenAsrSession * cc, valid[b] = 0; continue; } + // Same input-length gate as run(), before the row pads every other + // row's encoder input to its length. + if (timings[b].T_enc > T_enc_cap) { + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "qwen3_asr run_batch: utterance %d input too long — %d audio tokens, " + "at most %d fit the context with the prompt and output. See " + "transcribe_capabilities.max_audio_ms.", + b, timings[b].T_enc, T_enc_cap); + valid[b] = 0; + fail_status[b] = TRANSCRIBE_ERR_INPUT_TOO_LONG; + continue; + } n_chunks_max = std::max(n_chunks_max, timings[b].n_chunks); T_enc_out[b] = timings[b].T_enc; any = true; @@ -1219,16 +1332,9 @@ transcribe_status encode_all_batched(QwenAsrSession * cc, return TRANSCRIBE_OK; // caller emits per-row errors } - // Build batched encoder graph. if (const transcribe_status st = reset_compute_ctx(cc, 16); st != TRANSCRIBE_OK) { return st; } - EncoderBuildBatched eb = - build_encoder_graph_batched(cc->compute_ctx, cm->weights, cm->hparams, n_chunks_max, n, cc->encoder_use_flash); - if (eb.graph == nullptr || eb.out == nullptr) { - return TRANSCRIBE_ERR_GGUF; - } - if (cc->sched == nullptr) { cc->sched = ggml_backend_sched_new(cm->plan.scheduler_list.data(), nullptr, static_cast(cm->plan.scheduler_list.size()), 16384, false, true); @@ -1236,19 +1342,14 @@ transcribe_status encode_all_batched(QwenAsrSession * cc, return TRANSCRIBE_ERR_BACKEND; } } - ggml_backend_sched_reset(cc->sched); - if (!TRANSCRIBE_ALLOC_SCHED_GRAPH(cc->sched, eb.graph, "qwen3_asr")) { - return TRANSCRIBE_ERR_OOM; - } - - const int mel_per_chunk = cm->hparams.enc_n_window * 2; - const int T_per_chunk = eb.T_per_chunk; - const int T_pad_max = eb.T_pad_max; - const size_t per_chunk_elems = static_cast(mel_per_chunk) * n_mels; - // Pack mel: [mel_per_chunk, n_mels, 1, n*n_chunks_max]; utterance b's - // chunk c at N = b*n_chunks_max + c, zero-padded beyond n_chunks[b]. + // Conv subsampler over every packed chunk: [mel_per_chunk, n_mels, 1, + // n*n_chunks_max]; utterance b's chunk c at N = b*n_chunks_max + c, + // zero-padded beyond n_chunks[b]. + const int mel_per_chunk = cm->hparams.enc_n_window * 2; + std::vector subsample_host; { + const size_t per_chunk_elems = static_cast(mel_per_chunk) * n_mels; std::vector packed(per_chunk_elems * static_cast(n) * n_chunks_max, 0.0f); for (int b = 0; b < n; ++b) { if (!valid[b]) { @@ -1266,9 +1367,30 @@ transcribe_status encode_all_batched(QwenAsrSession * cc, } } } - ggml_backend_tensor_set(eb.mel_in, packed.data(), 0, packed.size() * sizeof(float)); + const int64_t t_sub0 = ggml_time_us(); + if (const transcribe_status st = run_subsample(cc, cm, packed.data(), n * n_chunks_max, subsample_host); + st != TRANSCRIBE_OK) { + return st; + } + enc_us += ggml_time_us() - t_sub0; } + // Build batched encoder graph. + EncoderBuildBatched eb = + build_encoder_graph_batched(cc->compute_ctx, cm->weights, cm->hparams, n_chunks_max, n, cc->encoder_use_flash); + if (eb.graph == nullptr || eb.out == nullptr) { + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_sched_reset(cc->sched); + if (!TRANSCRIBE_ALLOC_SCHED_GRAPH(cc->sched, eb.graph, "qwen3_asr")) { + return TRANSCRIBE_ERR_OOM; + } + + const int T_per_chunk = eb.T_per_chunk; + const int T_pad_max = eb.T_pad_max; + + ggml_backend_tensor_set(eb.subsample_in, subsample_host.data(), 0, subsample_host.size() * sizeof(float)); + // Positional embedding (shared across chunks/utterances). { std::vector pe = build_sinusoid_pe(cm->hparams.enc_d_model, T_per_chunk); @@ -1578,13 +1700,18 @@ transcribe_status run_batch(transcribe_session * session, std::vector valid(n, 0); // utterance produced a usable prefill int64_t mel_us = 0, enc_us = 0; + // Per-utterance terminal status for rejected rows. Defaults to INVALID_ARG; + // over-length rows are upgraded to INPUT_TOO_LONG. + std::vector fail_status(n, TRANSCRIBE_ERR_INVALID_ARG); + // Encode every utterance first so we can size the batched cache to the // real max prompt length before allocating it. std::vector> enc_hosts(n); std::vector T_enc(n, 0); + const int T_enc_cap = ceiling - k_gen_reserve - prompt_tokens_without_system(0, lang_prefix_ptr); const int64_t t_encpass0 = ggml_time_us(); - if (const transcribe_status st = - encode_all_batched(cc, cm, pcm, n_samples, n, valid, enc_hosts, T_enc, mel_us, enc_us); + if (const transcribe_status st = encode_all_batched(cc, cm, pcm, n_samples, n, T_enc_cap, valid, fail_status, + enc_hosts, T_enc, mel_us, enc_us); st != TRANSCRIBE_OK) { return st; } @@ -1595,9 +1722,6 @@ transcribe_status run_batch(transcribe_session * session, int max_T_prompt = 0; int max_T_enc = 0; int prefix_len = 0; - // Per-utterance terminal status for rejected rows. Defaults to INVALID_ARG; - // over-length rows below are upgraded to INPUT_TOO_LONG. - std::vector fail_status(n, TRANSCRIBE_ERR_INVALID_ARG); std::vector> prompt_ids(n); for (int b = 0; b < n; ++b) { if (!valid[b]) { From e6496d58dc199330b70f8bc6696e2170bdc646b1 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 22 Oct 2026 14:58:09 +0800 Subject: [PATCH 6/6] qwen3_asr, granite_nar: drop unused debug options and leftover mask plumbing Review follow-ups on top of b490c12f: - Remove TRANSCRIBE_QWEN3_ASR_CONV_PASS; the pass size is the constant kConvPassChunks. - Remove TRANSCRIBE_GRAPH_TOP (undocumented, unused outside manual debugging); transcribe-backend.{h,cpp} are back to main. - qwen3_asr: drop the always-null single-shot EncoderBuild::mask_in, and the orphaned build_cu_seqlens_mask doc comment and mask-upload note. - Tolerance notes: granite_nar decoder attention is flash by default and the encoder depthwise conv is the direct op; qwen3_asr encoder attention is flash by default. Re-measured on CPU (validate.py all, BF16): granite_nar 15/15, qwen3_asr 0.6B 13/13 within tolerance, transcripts match the reference. - clang-format the qwen3_asr conv-pass change and tests/granite_nar_bpe_pool_unit.cpp (the format gate failed on both). --- docs/environment-variables.md | 1 - src/arch/qwen3_asr/encoder.cpp | 7 ++--- src/arch/qwen3_asr/encoder.h | 7 +---- src/arch/qwen3_asr/model.cpp | 29 ++++--------------- src/transcribe-backend.cpp | 44 ----------------------------- src/transcribe-backend.h | 3 -- tests/granite_nar_bpe_pool_unit.cpp | 33 +++++++++++----------- tests/tolerances/granite_nar.json | 11 +++++--- tests/tolerances/qwen3_asr.json | 4 ++- 9 files changed, 34 insertions(+), 105 deletions(-) diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 04c2a233..9fa946e6 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -36,7 +36,6 @@ of tests. | `TRANSCRIBE_PERF_DEBUG` | Print a per-stage timing breakdown to stderr (DEBUG log) on the families that profile (`cohere`, `granite`, `canary`, `canary_qwen`, `moonshine`, `moonshine_streaming`, `moss`, `qwen3_asr`, `whisper`). For whisper, a value containing `cpu` or `all` additionally prints the CPU sub-section breakdown. | | `TRANSCRIBE_VOXTRAL_REALTIME_STREAM_TIMING` | Print a per-component streaming wall-time breakdown at stream finalize (voxtral_realtime). | | `TRANSCRIBE_MOONSHINE_STREAMING_ENC_WINDOW=` | Encoder frames emitted per window by the moonshine_streaming one-shot / batch encoder (default 1024, ~20 s). Windowing bounds encoder memory; it is mathematically exact (each window carries the full sliding-window context) and float-level differences stay at the backend's reduction-order noise. Set small to exercise the windowed path on short clips. Read at each run. | -| `TRANSCRIBE_QWEN3_ASR_CONV_PASS=` | Conv-subsampler chunks (1 s of audio each) per subsample-graph run in the qwen3_asr encoder (default 512, ~8.5 min). Longer inputs run the subsampler in passes so its graph stays a fixed size; chunks are independent, so the value does not change the result (bit-identical on CPU; GPU backends may differ at reduction-order noise for values that are not a multiple of 16). Set small to exercise the multi-pass path on short clips. Read at each run. | | `TRANSCRIBE_TEST_DEV_INIT_THROW=` | Fault injection: backend device init (`ggml_backend_dev_init`) throws for devices whose name contains `` (`*` matches every device). Exercises throw → skip → CPU-fallback in backend probing; an explicit backend request fails with `TRANSCRIBE_ERR_BACKEND`. Used by `backend_init_throw_unit` and `scripts/ci/vulkan_degradation_check.py`. | | `TRANSCRIBE_TEST_TEARDOWN_THROW` | Fault injection: any non-empty value injects a throw after each real free inside the `transcribe::safe_*` teardown wrappers, proving containment without leaking the handle. Used by `teardown_safety_unit`. | diff --git a/src/arch/qwen3_asr/encoder.cpp b/src/arch/qwen3_asr/encoder.cpp index f1ecc006..85afb7c7 100644 --- a/src/arch/qwen3_asr/encoder.cpp +++ b/src/arch/qwen3_asr/encoder.cpp @@ -18,7 +18,7 @@ namespace transcribe::qwen3_asr { -// Timing / cu_seqlens +// Timing int32_t aftercnn_len(int32_t mel_len) { // Three rounds of Conv2d(stride=2, pad=1, kernel=3): @@ -344,9 +344,6 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, named(eb.pos_emb_in, "enc.pos_emb.in"); ggml_set_input(eb.pos_emb_in); - // No attention mask: full unmasked attention (see encoder.h). - eb.mask_in = nullptr; - // Add sinusoidal PE (broadcast across n_chunks). pos_emb_in has // ne=[d_model, T_per_chunk]; ggml_add broadcasts missing axes. ggml_tensor * x = ggml_add(ctx, eb.subsample_in, eb.pos_emb_in); @@ -385,7 +382,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.dumps.pos_add_out = x; transcribe::debug::mark_tensor_for_dump(x); - // ----- 18 encoder blocks ----- + // ----- 18 encoder blocks (unmasked full attention, see encoder.h) ----- const int n_layers = static_cast(weights.enc_blocks.size()); for (int i = 0; i < n_layers; ++i) { x = build_enc_block(ctx, x, /*mask=*/nullptr, weights.enc_blocks[i], static_cast(d_model), diff --git a/src/arch/qwen3_asr/encoder.h b/src/arch/qwen3_asr/encoder.h index fbaec628..4d272cab 100644 --- a/src/arch/qwen3_asr/encoder.h +++ b/src/arch/qwen3_asr/encoder.h @@ -75,10 +75,6 @@ EncoderTiming compute_encoder_timing(int32_t n_mel_frames, const QwenAsrHParams // sin(p * inv_ts[k]), the second d_model/2 are cos(p * inv_ts[k]). std::vector build_sinusoid_pe(int32_t d_model, int32_t length, double max_timescale = 10000.0); -// Additive attention bias [T_enc, T_enc]. All zeros — full bidirectional -// attention over the pad-trimmed sequence (the eager baseline, not vLLM's -// cu_seqlens block-diagonal pattern). - // Conv chunks per group inside a subsample graph. The 3x Conv2d frontend is // independent per n_window*2-mel chunk, but its activations are large: // conv0's output is [50, 64, 480] f32 (6.1 MB) per chunk and conv1's im2col @@ -123,8 +119,7 @@ struct EncoderDumps { struct EncoderBuild { ggml_tensor * subsample_in = nullptr; // [d_model, per_chunk_aftercnn, n_chunks] ggml_tensor * pos_emb_in = nullptr; // [d_model, per_chunk_aftercnn] - ggml_tensor * mask_in = nullptr; // unused (null): full unmasked attention - ggml_tensor * out = nullptr; // [output_dim, T_enc] + ggml_tensor * out = nullptr; // [output_dim, T_enc] EncoderDumps dumps{}; ggml_cgraph * graph = nullptr; EncoderTiming timing{}; diff --git a/src/arch/qwen3_asr/model.cpp b/src/arch/qwen3_asr/model.cpp index bccb263e..db14e80f 100644 --- a/src/arch/qwen3_asr/model.cpp +++ b/src/arch/qwen3_asr/model.cpp @@ -563,23 +563,9 @@ void pack_mel_chunks(const float * mel, // [n_mels, T_mel] } } -// TRANSCRIBE_QWEN3_ASR_CONV_PASS= overrides kConvPassChunks, so the -// multi-pass path can be exercised (and checked against one pass) on short -// clips. Values below 1 are ignored. -int conv_pass_chunks() { - if (const char * v = transcribe::env::str("TRANSCRIBE_QWEN3_ASR_CONV_PASS"); v != nullptr) { - char * end = nullptr; - const long n = std::strtol(v, &end, 10); - if (end != v && n >= 1 && n <= (1 << 20)) { - return static_cast(n); - } - } - return kConvPassChunks; -} - // Conv subsampler over `n_chunks` packed mel chunks (pack_mel_chunks layout) // into out = [d_model, per_chunk_aftercnn, n_chunks], in passes of -// conv_pass_chunks() chunks. Full passes share one graph, allocated once and +// kConvPassChunks chunks. Full passes share one graph, allocated once and // re-run; the remainder gets one graph of its own size. Neither node count // nor scratch grows with the audio, and chunks are independent, so the // result matches one graph over every chunk. Built in cc->compute_ctx; @@ -619,10 +605,9 @@ transcribe_status run_subsample(QwenAsrSession * cc, return TRANSCRIBE_OK; }; - const int pass = conv_pass_chunks(); - const int n_full = n_chunks - n_chunks % pass; + const int n_full = n_chunks - n_chunks % kConvPassChunks; if (n_full > 0) { - if (const transcribe_status st = run_passes(pass, 0, n_full); st != TRANSCRIBE_OK) { + if (const transcribe_status st = run_passes(kConvPassChunks, 0, n_full); st != TRANSCRIBE_OK) { return st; } } @@ -796,10 +781,6 @@ transcribe_status run(transcribe_session * session, ggml_backend_tensor_set(eb.pos_emb_in, pe.data(), 0, pe.size() * sizeof(float)); } - // No attention mask: the encoder attends unmasked over the whole - // utterance (an all-zero additive mask is the same thing), so there is - // nothing to upload. - transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); const int64_t t_enc_start = ggml_time_us(); @@ -1708,8 +1689,8 @@ transcribe_status run_batch(transcribe_session * session, // real max prompt length before allocating it. std::vector> enc_hosts(n); std::vector T_enc(n, 0); - const int T_enc_cap = ceiling - k_gen_reserve - prompt_tokens_without_system(0, lang_prefix_ptr); - const int64_t t_encpass0 = ggml_time_us(); + const int T_enc_cap = ceiling - k_gen_reserve - prompt_tokens_without_system(0, lang_prefix_ptr); + const int64_t t_encpass0 = ggml_time_us(); if (const transcribe_status st = encode_all_batched(cc, cm, pcm, n_samples, n, T_enc_cap, valid, fail_status, enc_hosts, T_enc, mel_us, enc_us); st != TRANSCRIBE_OK) { diff --git a/src/transcribe-backend.cpp b/src/transcribe-backend.cpp index cf20cd31..41bd8071 100644 --- a/src/transcribe-backend.cpp +++ b/src/transcribe-backend.cpp @@ -7,10 +7,8 @@ #include "transcribe-backend.h" #include "ggml.h" -#include "transcribe-env.h" #include "transcribe-log.h" -#include #include #include #include @@ -203,47 +201,6 @@ void log_sched_scratch(ggml_backend_sched_t sched, const char * tag, int line, l } } -// TRANSCRIBE_GRAPH_TOP=N: log the N largest nodes of every graph allocated -// through alloc_sched_graph, with op, type, shape and assigned backend. The -// scratch line above says how big a graph is; this says which tensors make it -// that big. Sizes are per tensor, not the live set, so a tensor that is freed -// before the next big one is allocated still shows up. Debugging aid only. -void log_graph_top(ggml_backend_sched_t sched, struct ggml_cgraph * graph, const char * tag, int line) { - const char * env_n = env::str("TRANSCRIBE_GRAPH_TOP"); - if (env_n == nullptr) { - return; - } - const int top_n = std::atoi(env_n); - if (top_n <= 0) { - return; - } - std::vector nodes; - const int n_nodes = ggml_graph_n_nodes(graph); - nodes.reserve(static_cast(n_nodes)); - size_t total = 0; - for (int i = 0; i < n_nodes; ++i) { - ggml_tensor * t = ggml_graph_node(graph, i); - if (t->view_src != nullptr) { - continue; // views own no memory - } - total += ggml_nbytes(t); - nodes.push_back(t); - } - const size_t keep = std::min(nodes.size(), static_cast(top_n)); - std::partial_sort(nodes.begin(), nodes.begin() + static_cast(keep), nodes.end(), - [](const ggml_tensor * a, const ggml_tensor * b) { return ggml_nbytes(a) > ggml_nbytes(b); }); - log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, "%s@%d: graph top %zu of %d nodes (%.1f MB if nothing were reused)", tag, line, - keep, n_nodes, static_cast(total) / 1e6); - for (size_t i = 0; i < keep; ++i) { - const ggml_tensor * t = nodes[i]; - ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, const_cast(t)); - log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG, " %9.1f MB %-14s %-5s [%lld, %lld, %lld, %lld] %s %s", - static_cast(ggml_nbytes(t)) / 1e6, ggml_op_desc(t), ggml_type_name(t->type), - static_cast(t->ne[0]), static_cast(t->ne[1]), static_cast(t->ne[2]), - static_cast(t->ne[3]), backend != nullptr ? ggml_backend_name(backend) : "-", t->name); - } -} - } // namespace void safe_backend_free(ggml_backend_t backend) noexcept { @@ -287,7 +244,6 @@ bool alloc_sched_graph(ggml_backend_sched_t sched, return false; } log_sched_scratch(sched, tag, line, extent); - log_graph_top(sched, graph, tag, line); return true; } diff --git a/src/transcribe-backend.h b/src/transcribe-backend.h index 82efa359..f1c106ec 100644 --- a/src/transcribe-backend.h +++ b/src/transcribe-backend.h @@ -112,9 +112,6 @@ void release_compute_scratch(ggml_backend_sched_t & sched, struct ggml_context * // buffer's high-water mark, so per-token step reallocations stay silent. // `tag`/`line` identify the call site; `extent` (>= 0) is an optional // graph-length field to log alongside it (e.g. the encoder frame count). -// -// TRANSCRIBE_GRAPH_TOP=N additionally logs the N largest nodes of each graph -// (op, type, shape, backend, name), to attribute the scratch to tensors. bool alloc_sched_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph, const char * tag, diff --git a/tests/granite_nar_bpe_pool_unit.cpp b/tests/granite_nar_bpe_pool_unit.cpp index 7c65024e..7b1f778c 100644 --- a/tests/granite_nar_bpe_pool_unit.cpp +++ b/tests/granite_nar_bpe_pool_unit.cpp @@ -72,14 +72,14 @@ struct Head { // One scenario: `hidden` is [T, dim] row-major (frame-major), importance is // [T]. Runs the reference (frame-level logits, then pool + collapse) and the // graph's shape (pool hidden, then head + argmax + collapse) and compares. -void check_equivalence(const char * what, - const Head & head, +void check_equivalence(const char * what, + const Head & head, const std::vector & hidden, const std::vector & importance, int pool_window, int blank_id) { - const int T = static_cast(importance.size()); - const int n_win = (T + pool_window - 1) / pool_window; + const int T = static_cast(importance.size()); + const int n_win = (T + pool_window - 1) / pool_window; // --- Reference: frame-level logits, host pool + greedy (encoder.cpp) --- std::vector frame_logits; @@ -89,9 +89,9 @@ void check_equivalence(const char * what, blank_id, ref_ids); // --- Graph shape: importance-weighted hidden pool, then head + argmax --- - std::vector pooled(static_cast(n_win) * head.dim, 0.0f); - std::vector mass(static_cast(n_win), 0.0f); - std::vector window_ids(static_cast(n_win), blank_id); + std::vector pooled(static_cast(n_win) * head.dim, 0.0f); + std::vector mass(static_cast(n_win), 0.0f); + std::vector window_ids(static_cast(n_win), blank_id); for (int w = 0; w < n_win; ++w) { const int t0 = w * pool_window; const int t1 = std::min(t0 + pool_window, T); @@ -105,8 +105,7 @@ void check_equivalence(const char * what, for (int t = t0; t < t1; ++t) { const float wt = importance[static_cast(t)] / denom; for (int i = 0; i < head.dim; ++i) { - pooled[static_cast(w) * head.dim + i] += - wt * hidden[static_cast(t) * head.dim + i]; + pooled[static_cast(w) * head.dim + i] += wt * hidden[static_cast(t) * head.dim + i]; } } } @@ -159,8 +158,8 @@ int main() { Head head(/*dim=*/4, /*vocab=*/16); // T = 10, P = 4 -> 3 windows, the last one partial (2 frames). - const int T = 10; - const int P = 4; + const int T = 10; + const int P = 4; std::vector hidden(static_cast(T) * head.dim); for (int t = 0; t < T; ++t) { for (int i = 0; i < head.dim; ++i) { @@ -190,22 +189,22 @@ int main() { // Collapse rules: adjacent duplicate windows collapse, blanks are dropped. { - const std::vector ids = {3, 3, 0, 7, 0, 0, 7}; - const std::vector mass = {1, 1, 1, 1, 1, 1, 1}; + const std::vector ids = { 3, 3, 0, 7, 0, 0, 7 }; + const std::vector mass = { 1, 1, 1, 1, 1, 1, 1 }; std::vector out; CHECK(transcribe::granite_nar::collapse_bpe_window_ids(ids, mass, /*blank_id=*/0, out)); // Collapse runs on the window ids (before the shift): 3,3 -> one 3; // 0 is blank; 7 survives both times because a blank resets "prev". // Scheme A then subtracts 1 from each emitted id. - const std::vector expect = {2, 6, 6}; + const std::vector expect = { 2, 6, 6 }; CHECK(out == expect); } // A length mismatch is a caller bug: refuse instead of truncating. { - const std::vector ids = {1, 2, 3}; - const std::vector mass = {1, 1}; - std::vector out = {42}; + const std::vector ids = { 1, 2, 3 }; + const std::vector mass = { 1, 1 }; + std::vector out = { 42 }; CHECK(!transcribe::granite_nar::collapse_bpe_window_ids(ids, mass, /*blank_id=*/0, out)); CHECK(out.empty()); } diff --git a/tests/tolerances/granite_nar.json b/tests/tolerances/granite_nar.json index 3bd8106e..dbfc15df 100644 --- a/tests/tolerances/granite_nar.json +++ b/tests/tolerances/granite_nar.json @@ -11,8 +11,11 @@ " then processor.batch_decode(output.preds). LibriSpeech test-clean", " WER 1.29% (matches the model card).", " - C++: bf16 weights (F32 norms / biases, F16 conv kernels), F32", - " accumulators, CPU backend, threads=1. Bidirectional attention via", - " ggml_soft_max_ext(kq, mask=nullptr) in decoder.cpp.", + " accumulators, CPU backend, threads=1. Bidirectional decoder", + " attention via ggml_flash_attn_ext with no mask (the default;", + " TRANSCRIBE_NO_FLASH restores ggml_soft_max_ext(kq, mask=nullptr)).", + " Encoder depthwise conv via the direct op, not im2col + mul_mat.", + " Re-measured with both on: every tensor within the budgets below.", " - Mel: production C++ MelFrontend in eval mode against the HF", " GraniteSpeechNarFeatureExtractor (torchaudio MelSpectrogram +", " log10/max-8/div-4/add-1 + 2-frame stack).", @@ -82,7 +85,7 @@ "enc.block.0.post_conv": { "max_abs": 0.1, "mean_abs": 0.005, - "_comment": "Block 0 sub-step: post-conv residual. Observed 6.3e-2 / 3.3e-3." + "_comment": "Block 0 sub-step: post-conv residual. Observed 6.3e-2 / 3.3e-3 (6.2e-2 / 3.4e-3 with the direct depthwise conv)." }, "enc.block.0.post_ff2": { "max_abs": 0.1, @@ -132,6 +135,6 @@ "dec.text_logits": { "max_abs": 0.6, "mean_abs": 0.06, - "_comment": "Per-sample text-portion lm_head logits with /logits_scaling=8 applied (matching the 99a4df9 modeling — see decoder.cpp). 40 bf16 LLM layers + bf16 lm_head accumulate noise on post-scaling magnitudes of ~O(20). Observed 4.0e-1 / 3.5e-2." + "_comment": "Per-sample text-portion lm_head logits with /logits_scaling=8 applied (matching the 99a4df9 modeling — see decoder.cpp). 40 bf16 LLM layers + bf16 lm_head accumulate noise on post-scaling magnitudes of ~O(20). Observed 4.0e-1 / 3.5e-2 (4.1e-1 / 3.5e-2 with flash attention)." } } diff --git a/tests/tolerances/qwen3_asr.json b/tests/tolerances/qwen3_asr.json index 701c3971..28e54f85 100644 --- a/tests/tolerances/qwen3_asr.json +++ b/tests/tolerances/qwen3_asr.json @@ -5,7 +5,9 @@ "Reference: Alibaba qwen_asr v0.0.6 (author repo, transformers host)", "running BF16 inference on CPU.", "C++: ggml compute (CPU/Metal/Vulkan), weights stored as BF16 in the", - "accuracy GGUF.", + "accuracy GGUF. Encoder attention is ggml_flash_attn_ext with no mask", + "by default (TRANSCRIBE_NO_FLASH restores soft_max_ext); re-measured", + "on CPU with it on, all tensors stay within these budgets.", "", "Budgets are calibrated to ~1.5-2x the max observed across CPU,", "Metal, and Vulkan on JFK. Vulkan dominates most encoder tensors",