diff --git a/src/arch/canary_qwen/encoder.cpp b/src/arch/canary_qwen/encoder.cpp index b434a052..464020cf 100644 --- a/src/arch/canary_qwen/encoder.cpp +++ b/src/arch/canary_qwen/encoder.cpp @@ -29,7 +29,7 @@ namespace { namespace conf = transcribe::conformer; -bool detect_direct_dw_in_block(const char * /*backend*/) { +bool detect_direct_dw(const char * /*backend*/) { return conf::resolve_conv_direct("TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", /*backend_default=*/true); } @@ -109,8 +109,13 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const char * backend_name) { conf::ConvPolicy policy{}; policy.direct_pw = conf::detect_direct_pw(backend_name); - policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); - policy.direct_dw_in_pre_encode = false; + policy.direct_dw_in_block = detect_direct_dw(backend_name); + // Direct pre-encode depthwise too: the im2col form materialised a 2.3 GB + // [9, F/4, T/4, 256] f32 copy at 306 s, the graph's largest tensor. + policy.direct_dw_in_pre_encode = detect_direct_dw(backend_name); + // And window the stack in time (exact; see build_pre_encode): its + // whole-utterance activations are linear in audio length, 2 GB at 306 s. + policy.pre_encode_time_chunk = conf::resolve_pre_encode_time_chunk(conf::kPreEncodeTimeChunk); EncoderBuild eb{}; @@ -224,7 +229,12 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.out = x; ggml_set_output(eb.out); - eb.graph = ggml_new_graph_custom(ctx, 16384, false); + // 16384 for the blocks + perception projection; the pre_encode's time + // windows scale with audio length, so their nodes are added on top. + const size_t graph_size = + 16384 + static_cast(conf::kPreEncodeNodesPerWindow) * + static_cast(conf::pre_encode_window_count(n_mel_frames, policy.pre_encode_time_chunk)); + eb.graph = ggml_new_graph_custom(ctx, graph_size, false); if (eb.graph == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "canary_qwen encoder: ggml_new_graph_custom failed"); return eb; diff --git a/src/arch/canary_qwen/model.cpp b/src/arch/canary_qwen/model.cpp index 5b00b406..722e5461 100644 --- a/src/arch/canary_qwen/model.cpp +++ b/src/arch/canary_qwen/model.cpp @@ -19,6 +19,7 @@ #include "canary_qwen.h" #include "causal_lm/causal_lm.h" +#include "conformer/conformer.h" #include "decoder.h" #include "encoder.h" #include "ggml-alloc.h" @@ -802,7 +803,8 @@ transcribe_status run(transcribe_session * context, } { ggml_init_params ip{}; - ip.mem_size = 32 * 1024 * 1024; + ip.mem_size = 32 * 1024 * 1024 + static_cast(conformer::pre_encode_meta_bytes( + mel_n_frames, conformer::kPreEncodeTimeChunk)); ip.mem_buffer = nullptr; ip.no_alloc = true; cc->compute_ctx = ggml_init(ip); @@ -820,7 +822,8 @@ transcribe_status run(transcribe_session * context, if (cc->sched == nullptr) { cc->sched = ggml_backend_sched_new(cm->plan.scheduler_list.data(), nullptr, - static_cast(cm->plan.scheduler_list.size()), 16384, /*parallel=*/false, + static_cast(cm->plan.scheduler_list.size()), + 16384 + conformer::kPreEncodeExtraGraphNodes, /*parallel=*/false, /*op_offload=*/true); if (cc->sched == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "canary_qwen run: ggml_backend_sched_new failed"); @@ -1316,7 +1319,8 @@ transcribe_status encode_one(CanaryQwenSession * cc, 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); + static_cast(cm->plan.scheduler_list.size()), + 16384 + conformer::kPreEncodeExtraGraphNodes, false, true); if (cc->sched == nullptr) { return TRANSCRIBE_ERR_BACKEND; } diff --git a/src/arch/cohere/encoder.cpp b/src/arch/cohere/encoder.cpp index 0b0819cc..d2fd7400 100644 --- a/src/arch/cohere/encoder.cpp +++ b/src/arch/cohere/encoder.cpp @@ -23,13 +23,17 @@ namespace { namespace conf = transcribe::conformer; -// Depthwise dispatch policy: direct ggml_conv_2d_dw_direct on Vulkan/CUDA -// (native kernels), im2col + mul_mat on Metal/CPU. See the ConvPolicy -// comment in conformer.h for why the two dw sites share one detect result. -bool detect_direct_dw(const char * backend) { - const bool backend_default = - backend != nullptr && (std::strstr(backend, "Vulkan") != nullptr || std::strstr(backend, "CUDA") != nullptr); - return conf::resolve_conv_direct("TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", backend_default); +// Depthwise dispatch policy: the direct CONV_2D_DW op on every backend, for +// both depthwise sites. Metal and CPU used to take im2col + mul_mat, which at +// 306 s materialised a 2.3 GB [9, F/4, T/4, 256] f32 copy for the first +// pre-encode depthwise conv alone (the graph's largest tensor), and the block +// site switched algorithm at B > 1 (conv_1d_dw_f32), so a batched utterance +// was not computed the way its single-shot self was. CONV_2D_DW is tested on +// every device present by tests/conv_dw_direct_unit.cpp. +// TRANSCRIBE_CONV_NO_DIRECT_DW=1 restores im2col (single-utterance only). +bool detect_direct_dw(const char * /*backend*/) { + return conf::resolve_conv_direct("TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", + /*backend_default=*/true); } // ----- Views ------------------------------------------------------ @@ -116,6 +120,10 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const bool direct_dw = detect_direct_dw(backend_name); policy.direct_dw_in_block = direct_dw; policy.direct_dw_in_pre_encode = direct_dw; + // Window the pre_encode stack in time: its whole-utterance activations + // were 3 GB at 306 s (conv0 output, its matmul twin, ReLU), linear in + // audio length. Exact; see build_pre_encode. + policy.pre_encode_time_chunk = conf::resolve_pre_encode_time_chunk(conf::kPreEncodeTimeChunk); EncoderBuild eb{}; @@ -222,8 +230,12 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, ggml_set_output(eb.out); // Build the forward cgraph. 48 blocks + enc_dec_proj means a - // large graph; 16384 leaves headroom. - eb.graph = ggml_new_graph_custom(ctx, 16384, false); + // large graph; 16384 leaves headroom for those. The pre_encode's time + // windows scale with audio length, so their nodes are added on top. + const size_t graph_size = + 16384 + static_cast(conf::kPreEncodeNodesPerWindow) * + static_cast(conf::pre_encode_window_count(n_mel_frames, policy.pre_encode_time_chunk)); + eb.graph = ggml_new_graph_custom(ctx, graph_size, false); if (eb.graph == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "cohere encoder: ggml_new_graph_custom failed"); return eb; diff --git a/src/arch/cohere/model.cpp b/src/arch/cohere/model.cpp index e7f9e054..db107fde 100644 --- a/src/arch/cohere/model.cpp +++ b/src/arch/cohere/model.cpp @@ -6,6 +6,7 @@ // on the ggml graph (not host-side like Parakeet's LSTM). #include "cohere.h" +#include "conformer/conformer.h" #include "decoder.h" #include "encoder.h" #include "ggml-alloc.h" @@ -713,9 +714,11 @@ transcribe_status run(transcribe_session * session, // ----- Build encoder graph ------------------------------------- { - // 48 encoder blocks + decoder = large graph. 8 MB metadata arena. + // 48 encoder blocks + decoder = large graph. 8 MB metadata arena, plus + // the pre_encode's per-window tensors (it tiles in time). ggml_init_params init_params{}; - init_params.mem_size = 8 * 1024 * 1024; + init_params.mem_size = + 8 * 1024 * 1024 + conformer::pre_encode_meta_bytes(mel_n_frames, conformer::kPreEncodeTimeChunk); init_params.mem_buffer = nullptr; init_params.no_alloc = true; cc->compute_ctx = ggml_init(init_params); @@ -743,7 +746,8 @@ transcribe_status run(transcribe_session * session, // ----- 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); + static_cast(cm->plan.scheduler_list.size()), + 16384 + conformer::kPreEncodeExtraGraphNodes, false, true); if (cc->sched == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "cohere run: ggml_backend_sched_new failed"); return TRANSCRIBE_ERR_BACKEND; @@ -1408,7 +1412,8 @@ transcribe_status encode_one_to_host(CohereSession * cc, 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); + static_cast(cm->plan.scheduler_list.size()), + 16384 + conformer::kPreEncodeExtraGraphNodes, false, true); if (cc->sched == nullptr) { return TRANSCRIBE_ERR_BACKEND; } diff --git a/src/arch/granite/encoder.cpp b/src/arch/granite/encoder.cpp index 1396df66..4030be55 100644 --- a/src/arch/granite/encoder.cpp +++ b/src/arch/granite/encoder.cpp @@ -252,9 +252,23 @@ ggml_tensor * granite_conv_module(ggml_context * ctx, // Depthwise conv1d: kernel [k, 1, inner_dim], symmetric pad (k-1)/2 // (= 7 for k=15). Granite's depthwise has no bias. + // + // 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 by 7x) + // to feed a matmul whose K is only conv_kernel. Same swap as + // granite_nar and granite5_ctc; not bitwise identical (accumulation + // order), so it is gated on transcripts. 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); + { + 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); + } // Fused BatchNorm: y = x * fused_scale + fused_bias, with 1-D scale // and bias broadcast across the time axis. diff --git a/src/arch/parakeet/encoder.cpp b/src/arch/parakeet/encoder.cpp index ba0e6253..c608220d 100644 --- a/src/arch/parakeet/encoder.cpp +++ b/src/arch/parakeet/encoder.cpp @@ -55,16 +55,16 @@ bool detect_direct_dw_in_block(const char * backend) { /*backend_default=*/true); } -// Pre-encode depthwise dispatch: direct path EXCEPT on Metal, where -// CONV_2D_DW is unsupported in this vendored ggml (see conv_2d_dw_f32 in -// conformer.cpp) so the im2col helper is used. CPU/Vulkan have native -// CONV_2D_DW; the direct path avoids the ~31x im2col expansion and the -// degenerate per-channel [1 x K] matmul (~40-50ms here). +// Pre-encode depthwise dispatch: the direct path on every backend; it avoids +// the ~31x im2col expansion and the degenerate per-channel [1 x K] matmul +// (~40-50ms here). This used to exclude Metal on the belief that CONV_2D_DW +// was unsupported there, but it matched on the string "Metal" and the Metal +// backend is named "MTL0", so Metal has always taken the direct path; it is +// supported and tested (tests/conv_dw_direct_unit.cpp). bool detect_direct_dw_in_pre_encode(const char * backend) { - const bool is_metal = - backend != nullptr && (std::strstr(backend, "Metal") != nullptr || std::strstr(backend, "metal") != nullptr); + (void) backend; return conf::resolve_conv_direct("TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", - /*backend_default=*/!is_metal); + /*backend_default=*/true); } // ----- Views ------------------------------------------------------ diff --git a/src/arch/sortformer/model.cpp b/src/arch/sortformer/model.cpp index 09eb4455..6f9c0345 100644 --- a/src/arch/sortformer/model.cpp +++ b/src/arch/sortformer/model.cpp @@ -326,10 +326,11 @@ conf::ConvPolicy sf_conv_policy(const char * backend) { // Parakeet block/pre_encode depthwise defaults (encoder.cpp:50-66). policy.direct_dw_in_block = conf::resolve_conv_direct("TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", /*backend_default=*/true); - const bool is_metal = - backend != nullptr && (std::strstr(backend, "Metal") != nullptr || std::strstr(backend, "metal") != nullptr); + // Direct on every backend, Metal included (parakeet's + // detect_direct_dw_in_pre_encode explains the old Metal exclusion, which + // never matched the "MTL0" backend name). policy.direct_dw_in_pre_encode = conf::resolve_conv_direct( - "TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", /*backend_default=*/!is_metal); + "TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", /*backend_default=*/true); policy.causal_pre_encode = false; // NEST offline subsample (Regular att) return policy; } diff --git a/src/conformer/conformer.cpp b/src/conformer/conformer.cpp index 9b6c2ded..6df044f3 100644 --- a/src/conformer/conformer.cpp +++ b/src/conformer/conformer.cpp @@ -13,6 +13,7 @@ #include "transcribe-env.h" #include "transcribe-log.h" +#include #include #include #include @@ -179,9 +180,13 @@ ggml_tensor * conv_1d_f32(ggml_context * ctx, // kernel's real type to im2col). CANONICAL Metal conv-quirk note: vendored // ggml's ggml_conv_2d_dw hardcodes F16 im2col, which routes the matmul to // a `kernel_mul_mv_f32_f16_short` Metal kernel that isn't compiled in; -// forcing f32 lands on the available `kernel_mul_mv_f32_f32_short`. The -// alternative ggml_conv_2d_dw_direct can't help: its CONV_2D_DW op is -// unsupported on Metal, so the im2col + mul_mat path here is what works. +// forcing f32 lands on the available `kernel_mul_mv_f32_f32_short`. +// +// Prefer conv_2d_dw_direct_f32 below for new code: CONV_2D_DW runs on Metal, +// CPU, Vulkan and CUDA (tests/conv_dw_direct_unit.cpp checks it on every +// device present), takes one algorithm at every batch size, and does not +// materialise the kernel-width-fold copy this im2col form does (2.3 GB for +// the conformer pre-encode's first depthwise conv at 306 s). ggml_tensor * conv_2d_dw_f32(ggml_context * ctx, ggml_tensor * kernel, // [KW, KH, 1, C] ggml_tensor * data, // [W, H, C, N] @@ -211,8 +216,8 @@ ggml_tensor * conv_2d_dw_f32(ggml_context * ctx, // BF16-reference conformer (e.g. cohere) ships its depthwise kernels at F16, // so promote to F32 in-graph before any direct depthwise op. The depthwise // kernel is tiny, so the cast is negligible; F32 kernels pass through -// untouched. Only the direct_dw path (CUDA/Vulkan) routes here — the im2col -// fallback used on Metal/CPU takes the kernel's real type and is unaffected. +// untouched. Only the direct depthwise paths route here; the im2col forms +// take the kernel's real type and are unaffected. static ggml_tensor * dw_kernel_for_direct(ggml_context * ctx, ggml_tensor * kernel) { if (kernel == nullptr || kernel->type == GGML_TYPE_F32) { return kernel; @@ -337,6 +342,38 @@ bool detect_direct_pw(const char * backend) { return resolve_conv_direct("TRANSCRIBE_CONV_DIRECT_PW", "TRANSCRIBE_CONV_NO_DIRECT_PW", backend_default); } +int pre_encode_window_count(int64_t n_mel_frames, int chunk) { + if (chunk <= 0 || n_mel_frames <= 0) { + return 1; + } + auto stage_len = [](int64_t L) { + return (L + 2 * 1 - 3) / 2 + 1; + }; + const int64_t t_enc = stage_len(stage_len(stage_len(n_mel_frames))); + if (t_enc <= chunk) { + return 1; + } + // Matches the balanced split in build_pre_encode: ceil(T_enc / chunk) + // windows, capped (see kPreEncodeMaxWindows). + const int64_t n_win = (t_enc + chunk - 1) / chunk; + return static_cast(std::min(n_win, kPreEncodeMaxWindows)); +} + +size_t pre_encode_meta_bytes(int64_t n_mel_frames, int chunk) { + constexpr size_t kBytesPerWindow = 256 * 1024; + const int n_win = pre_encode_window_count(n_mel_frames, chunk); + return static_cast(n_win > 1 ? n_win - 1 : 0) * kBytesPerWindow; +} + +int resolve_pre_encode_time_chunk(int family_default) { + const char * v = transcribe::env::str("TRANSCRIBE_PRE_ENCODE_CHUNK"); + if (v == nullptr || *v == '\0') { + return family_default; + } + const long n = std::strtol(v, nullptr, 10); + return n > 0 ? static_cast(n) : 0; +} + // Conv module (pointwise -> GLU -> depthwise -> BN/LN -> SiLU -> pointwise). // Operates on the post-LayerNorm activation (LN applied by the caller). // Pointwise convs (k=1) are direct mul_mats in [d_model, T] layout by @@ -496,7 +533,7 @@ ggml_tensor * conv_module(ggml_context * ctx, ggml_tensor * x, const BlockView & // Output: [T_out, 1, d_model, B] → [T_out, d_model, B]. x = ggml_reshape_3d(ctx, x, x->ne[0], x->ne[2], x->ne[3]); } else { - // im2col depthwise (cohere on Metal/CPU); single-utterance only. + // im2col depthwise (TRANSCRIBE_CONV_NO_DIRECT_DW); single-utterance only. x = conv_1d_dw_f32(ctx, b.conv_dw_w, x, /*s=*/1, /*p=*/padding_op, /*d=*/1); } @@ -1225,144 +1262,234 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, return ggml_mul_inplace(ctx, t, m); }; - // conv0 (standard 2D conv: 1 in, channels out, k=3 s=2) - if (policy.direct_conv0_in_pre_encode) { - x = ggml_conv_2d_direct(ctx, pe.conv0_w, x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); - } else { - x = ggml_conv_2d(ctx, pe.conv0_w, x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); - } - x = add_conv_bias(ctx, x, pe.conv0_b); - x = name_prefixed(x, name_prefix, "conv0"); - x = ggml_relu(ctx, x); - x = name_prefixed(x, name_prefix, "relu0"); - x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s1 : nullptr, "pre_encode.valid_mask.s1"); - x = zero_causal_tail(x); - - // conv2 (depthwise: channels -> channels, groups=channels, k=3 s=2). - // im2col path (conv_2d_dw_f32) when direct_dw_in_pre_encode is false. - if (policy.direct_dw_in_pre_encode) { - x = ggml_conv_2d_dw_direct(ctx, dw_kernel_for_direct(ctx, pe.conv2_w), x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); - } else { - x = conv_2d_dw_f32(ctx, pe.conv2_w, x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); - } - x = add_conv_bias(ctx, x, pe.conv2_b); - x = name_prefixed(x, name_prefix, "conv2"); - const bool after_stride = policy.pre_encode_mask_after_stride; - if (after_stride) { - x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s2 : nullptr, "pre_encode.valid_mask.s2"); - } + // The conv stack proper: [F, H, 1, B] mel slice -> [F', H', C, B]. A + // lambda so the time-tiled path below can run it once per chunk; the + // masked and causal paths (which carry whole-utterance state) only ever + // run it once. + auto conv_stack = [&](ggml_tensor * x) -> ggml_tensor * { + // conv0 (standard 2D conv: 1 in, channels out, k=3 s=2) + if (policy.direct_conv0_in_pre_encode) { + x = ggml_conv_2d_direct(ctx, pe.conv0_w, x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } else { + x = ggml_conv_2d(ctx, pe.conv0_w, x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } + x = add_conv_bias(ctx, x, pe.conv0_b); + x = name_prefixed(x, name_prefix, "conv0"); + x = ggml_relu(ctx, x); + x = name_prefixed(x, name_prefix, "relu0"); + x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s1 : nullptr, "pre_encode.valid_mask.s1"); + x = zero_causal_tail(x); + + // conv2 (depthwise: channels -> channels, groups=channels, k=3 s=2). + // im2col path (conv_2d_dw_f32) when direct_dw_in_pre_encode is false. + if (policy.direct_dw_in_pre_encode) { + x = ggml_conv_2d_dw_direct(ctx, dw_kernel_for_direct(ctx, pe.conv2_w), x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } else { + x = conv_2d_dw_f32(ctx, pe.conv2_w, x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } + x = add_conv_bias(ctx, x, pe.conv2_b); + x = name_prefixed(x, name_prefix, "conv2"); + const bool after_stride = policy.pre_encode_mask_after_stride; + if (after_stride) { + x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s2 : nullptr, "pre_encode.valid_mask.s2"); + } - // conv3 (pointwise: channels -> channels, k=1 s=1 p=0) - x = policy.pre_encode_f32_pointwise ? - conv_2d_f32(ctx, pe.conv3_w, x, /*s0=*/1, /*s1=*/1, /*p0=*/0, /*p1=*/0, /*d0=*/1, /*d1=*/1) : - ggml_conv_2d(ctx, pe.conv3_w, x, - /*s0=*/1, /*s1=*/1, - /*p0=*/0, /*p1=*/0, - /*d0=*/1, /*d1=*/1); - x = add_conv_bias(ctx, x, pe.conv3_b); - x = name_prefixed(x, name_prefix, "conv3"); - x = ggml_relu(ctx, x); - x = name_prefixed(x, name_prefix, "relu3"); - if (!after_stride) { - x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s2 : nullptr, "pre_encode.valid_mask.s2"); - } else { - x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s2_extent : nullptr, "pre_encode.extent_mask.s2"); - } - x = zero_causal_tail(x); - - // conv5 (depthwise) -> conv6 (pointwise) -> ReLU - if (policy.direct_dw_in_pre_encode) { - x = ggml_conv_2d_dw_direct(ctx, dw_kernel_for_direct(ctx, pe.conv5_w), x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); - } else { - x = conv_2d_dw_f32(ctx, pe.conv5_w, x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); - } - x = add_conv_bias(ctx, x, pe.conv5_b); - x = name_prefixed(x, name_prefix, "conv5"); - if (after_stride) { - x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s3 : nullptr, "pre_encode.valid_mask.s3"); - } + // conv3 (pointwise: channels -> channels, k=1 s=1 p=0) + x = policy.pre_encode_f32_pointwise ? + conv_2d_f32(ctx, pe.conv3_w, x, /*s0=*/1, /*s1=*/1, /*p0=*/0, /*p1=*/0, /*d0=*/1, /*d1=*/1) : + ggml_conv_2d(ctx, pe.conv3_w, x, + /*s0=*/1, /*s1=*/1, + /*p0=*/0, /*p1=*/0, + /*d0=*/1, /*d1=*/1); + x = add_conv_bias(ctx, x, pe.conv3_b); + x = name_prefixed(x, name_prefix, "conv3"); + x = ggml_relu(ctx, x); + x = name_prefixed(x, name_prefix, "relu3"); + if (!after_stride) { + x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s2 : nullptr, "pre_encode.valid_mask.s2"); + } else { + x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s2_extent : nullptr, "pre_encode.extent_mask.s2"); + } + x = zero_causal_tail(x); + + // conv5 (depthwise) -> conv6 (pointwise) -> ReLU + if (policy.direct_dw_in_pre_encode) { + x = ggml_conv_2d_dw_direct(ctx, dw_kernel_for_direct(ctx, pe.conv5_w), x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } else { + x = conv_2d_dw_f32(ctx, pe.conv5_w, x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } + x = add_conv_bias(ctx, x, pe.conv5_b); + x = name_prefixed(x, name_prefix, "conv5"); + if (after_stride) { + x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s3 : nullptr, "pre_encode.valid_mask.s3"); + } - x = policy.pre_encode_f32_pointwise ? - conv_2d_f32(ctx, pe.conv6_w, x, /*s0=*/1, /*s1=*/1, /*p0=*/0, /*p1=*/0, /*d0=*/1, /*d1=*/1) : - ggml_conv_2d(ctx, pe.conv6_w, x, - /*s0=*/1, /*s1=*/1, - /*p0=*/0, /*p1=*/0, - /*d0=*/1, /*d1=*/1); - x = add_conv_bias(ctx, x, pe.conv6_b); - x = name_prefixed(x, name_prefix, "conv6"); - x = ggml_relu(ctx, x); - x = name_prefixed(x, name_prefix, "relu6"); - if (!after_stride) { - x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s3 : nullptr, "pre_encode.valid_mask.s3"); - } else { - x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s3_extent : nullptr, "pre_encode.extent_mask.s3"); - } - // Causal: drop the spurious tail of the last stage as a view; the - // flatten's cont below materializes the exact [F', T_enc] extent. - if (causal_pe) { - want_w = want_w / 2 + 1; - want_h = want_h / 2 + 1; - assert(x->ne[0] >= want_w && x->ne[0] - want_w <= 1 && x->ne[1] >= want_h && x->ne[1] - want_h <= 1); - if (x->ne[0] != want_w || x->ne[1] != want_h) { - x = ggml_view_4d(ctx, x, want_w, want_h, x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], /*offset=*/0); + x = policy.pre_encode_f32_pointwise ? + conv_2d_f32(ctx, pe.conv6_w, x, /*s0=*/1, /*s1=*/1, /*p0=*/0, /*p1=*/0, /*d0=*/1, /*d1=*/1) : + ggml_conv_2d(ctx, pe.conv6_w, x, + /*s0=*/1, /*s1=*/1, + /*p0=*/0, /*p1=*/0, + /*d0=*/1, /*d1=*/1); + x = add_conv_bias(ctx, x, pe.conv6_b); + x = name_prefixed(x, name_prefix, "conv6"); + x = ggml_relu(ctx, x); + x = name_prefixed(x, name_prefix, "relu6"); + if (!after_stride) { + x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s3 : nullptr, "pre_encode.valid_mask.s3"); + } else { + x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s3_extent : nullptr, "pre_encode.extent_mask.s3"); } - } + // Causal: drop the spurious tail of the last stage as a view; the + // flatten's cont below materializes the exact [F', T_enc] extent. + if (causal_pe) { + want_w = want_w / 2 + 1; + want_h = want_h / 2 + 1; + assert(x->ne[0] >= want_w && x->ne[0] - want_w <= 1 && x->ne[1] >= want_h && x->ne[1] - want_h <= 1); + if (x->ne[0] != want_w || x->ne[1] != want_h) { + x = ggml_view_4d(ctx, x, want_w, want_h, x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], + /*offset=*/0); + } + } + return x; + }; - // At this point ne = [F'=16, T_enc, channels=256, 1] where - // T_enc = floor(T_mel / 8). Flatten (F', C) into one feature axis - // matching pre_encode_in = channels * (n_mels / subsampling_factor). - // The reference uses flat index = c*F' + f' (C slow, F' fast), which in - // ggml ne is: permute F' to axis 0, C to axis 1, then collapse 0 and 1: - // permute (0, 2, 1, 3): [F', T, C, 1] -> [F', C, T, 1] - // ggml_cont; reshape -> [F'*C, T] = [pre_encode_in, T_enc] - const int64_t F_prime = x->ne[0]; - const int64_t T_enc = x->ne[1]; - const int64_t C = x->ne[2]; - // Utterance batch rides the conv "N" axis (ne[3]) through the stem; the - // flatten + linear below land it on ne[2] == B for the conformer blocks. - const int64_t B = x->ne[3]; - const int64_t pre_encode_in = F_prime * C; - - // Sanity check: the linear layer expects exactly pre_encode_in features. - if (pe.out_w != nullptr && pre_encode_in != pe.out_w->ne[0]) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "%s encoder: pre_encode_in mismatch: " - "F'*C=%lld but out_w expects %lld", - error_tag != nullptr ? error_tag : "conformer", static_cast(pre_encode_in), - static_cast(pe.out_w->ne[0])); - return nullptr; - } + // Flatten + linear, without the output bias: [F', T, C, B] -> [d_model, + // T, B]. Frame-local, so the tiled path applies it per chunk. + auto flatten_linear = [&](ggml_tensor * x) -> ggml_tensor * { + // At this point ne = [F'=16, T_enc, channels=256, 1] where + // T_enc = floor(T_mel / 8). Flatten (F', C) into one feature axis + // matching pre_encode_in = channels * (n_mels / subsampling_factor). + // The reference uses flat index = c*F' + f' (C slow, F' fast), which in + // ggml ne is: permute F' to axis 0, C to axis 1, then collapse 0 and 1: + // permute (0, 2, 1, 3): [F', T, C, 1] -> [F', C, T, 1] + // ggml_cont; reshape -> [F'*C, T] = [pre_encode_in, T_enc] + const int64_t F_prime = x->ne[0]; + const int64_t T_enc = x->ne[1]; + const int64_t C = x->ne[2]; + // Utterance batch rides the conv "N" axis (ne[3]) through the stem; the + // flatten + linear below land it on ne[2] == B for the conformer blocks. + const int64_t B = x->ne[3]; + const int64_t pre_encode_in = F_prime * C; + + // Sanity check: the linear layer expects exactly pre_encode_in features. + if (pe.out_w != nullptr && pre_encode_in != pe.out_w->ne[0]) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "%s encoder: pre_encode_in mismatch: " + "F'*C=%lld but out_w expects %lld", + error_tag != nullptr ? error_tag : "conformer", static_cast(pre_encode_in), + static_cast(pe.out_w->ne[0])); + return nullptr; + } - x = ggml_permute(ctx, x, /*axis0=*/0, /*axis1=*/2, - /*axis2=*/1, /*axis3=*/3); - x = name_prefixed(x, name_prefix, "permuted"); - x = ggml_cont(ctx, x); - x = name_prefixed(x, name_prefix, "flat_pre"); - // [F', C, T_enc, B] -> [F'*C, T_enc, B]. B at ne[2] after the collapse. - x = ggml_reshape_3d(ctx, x, pre_encode_in, T_enc, B); - x = name_prefixed(x, name_prefix, "flat"); - - // Linear projection to d_model (out_w ne=[pre_encode_in, d_model]). - x = ggml_mul_mat(ctx, pe.out_w, x); - x = name_prefixed(x, name_prefix, "linear"); + x = ggml_permute(ctx, x, /*axis0=*/0, /*axis1=*/2, + /*axis2=*/1, /*axis3=*/3); + x = name_prefixed(x, name_prefix, "permuted"); + x = ggml_cont(ctx, x); + x = name_prefixed(x, name_prefix, "flat_pre"); + // [F', C, T_enc, B] -> [F'*C, T_enc, B]. B at ne[2] after the collapse. + x = ggml_reshape_3d(ctx, x, pre_encode_in, T_enc, B); + x = name_prefixed(x, name_prefix, "flat"); + + // Linear projection to d_model (out_w ne=[pre_encode_in, d_model]). + x = ggml_mul_mat(ctx, pe.out_w, x); + x = name_prefixed(x, name_prefix, "linear"); + return x; + }; + + // Time tiling (policy.pre_encode_time_chunk > 0). The whole-utterance + // stack keeps conv0's [F/2, T/2, C] f32 output alive (with its im2col + // matmul twin) next to the first depthwise conv's input: 2 GB at 306 s + // for C=256, linear in audio length. Instead run the stack over + // overlapping mel windows and keep each window's exact output frames. + // + // Exactness. With k=3, s=2, p=1 at every stage, output frame o reads + // mel rows [8o-7, 8o+7]. A window starting at mel row m0 = 8(a-1) + // (a multiple of 8, so every stage's local grid lands on the global + // one) computes global frame a-1+j as its local frame j; local frame 0 + // reads the window's left zero pad where the utterance has real rows, + // so it is dropped, and local frames >= 1 read only rows inside the + // window. Ending the window at mel row 8b makes frame b-1 (rows up to + // 8b-1) the last exact one. The first window starts at row 0 and the + // last ends at T_mel, so the utterance edges see the same zero padding + // as the untiled graph. Every kept frame is computed from the same + // inputs by the same per-element ops, so the result is the untiled one. + // Bitwise, it matches wherever the backend's matmul kernel choice does + // not depend on the extent: Metal does (tests/pre_encode_tile_unit.cpp, + // and cohere's encoder dumps at 11-306 s) as long as no window drops to + // its few-row kernels (<= 16 columns), hence balanced windows, equal to + // within one frame and never shorter than half of `chunk`. CPU picks its + // matmul path by extent, so there the vendored F16-im2col pointwise + // convs differ by accumulation order (~1e-4 relative, F16 territory). + // + // Masked (batched variable-length) and causal stacks keep the + // whole-utterance path: their masks and causal tail are indexed by + // global length. + const int64_t T_mel = x->ne[1]; + auto stage_len = [](int64_t L) { + return (L + 2 * 1 - 3) / 2 + 1; + }; + const int64_t T_enc_out = stage_len(stage_len(stage_len(T_mel))); + const int64_t chunk = policy.pre_encode_time_chunk; + const bool tiled = chunk > 0 && valid_masks == nullptr && !causal_pe && T_enc_out > chunk; + if (!tiled) { + x = conv_stack(x); + x = (x != nullptr) ? flatten_linear(x) : nullptr; + if (x == nullptr) { + return nullptr; + } + } else { + ggml_tensor * mel_t = x; + ggml_tensor * out = nullptr; + const int64_t n_win = pre_encode_window_count(T_mel, static_cast(chunk)); + const int64_t base = T_enc_out / n_win; + const int64_t extra = T_enc_out % n_win; + for (int64_t w = 0, a = 0; w < n_win; ++w) { + const int64_t b = a + base + (w < extra ? 1 : 0); + const int64_t skip = (a == 0) ? 0 : 1; + const int64_t m0 = 8 * (a - skip); + const int64_t m1 = (b == T_enc_out) ? T_mel : 8 * b; + ggml_tensor * win = + ggml_view_4d(ctx, mel_t, mel_t->ne[0], m1 - m0, mel_t->ne[2], mel_t->ne[3], mel_t->nb[1], mel_t->nb[2], + mel_t->nb[3], static_cast(m0) * mel_t->nb[1]); + ggml_tensor * y = conv_stack(ggml_cont(ctx, win)); + if (y == nullptr) { + return nullptr; + } + assert(y->ne[1] == b - a + skip); + if (skip != 0) { + y = ggml_view_4d(ctx, y, y->ne[0], b - a, y->ne[2], y->ne[3], y->nb[1], y->nb[2], y->nb[3], + static_cast(skip) * y->nb[1]); + } + y = flatten_linear(y); + if (y == nullptr) { + return nullptr; + } + // Chaining the concat keeps the windows in dependency order, so + // the allocator retires one window's activations before the next. + out = (out == nullptr) ? y : ggml_concat(ctx, out, y, /*dim=*/1); + a = b; + } + assert(out->ne[1] == T_enc_out); + x = out; + } // Add bias [d_model], broadcast to [d_model, T_enc, 1, 1]. if (pe.out_b != nullptr) { diff --git a/src/conformer/conformer.h b/src/conformer/conformer.h index 74541634..202a3f0c 100644 --- a/src/conformer/conformer.h +++ b/src/conformer/conformer.h @@ -15,7 +15,8 @@ // debug-dump. Helpers name no intermediates except build_pre_encode (explicit // prefix). The f32-friendly conv helpers (conv_1d_f32, conv_2d_dw_f32, // conv_1d_dw_f32) work around ggml kernel limitations on Metal (see -// conv_2d_dw_f32 in conformer.cpp); use them for f32 kernels on Metal. +// conv_2d_dw_f32 in conformer.cpp); use them for f32 kernels on Metal. For +// depthwise convs prefer conv_2d_dw_direct_f32, which needs no workaround. #pragma once @@ -129,6 +130,13 @@ struct ConvPolicy { // Run the pre_encode pointwise convs (conv3, conv6) via conv_2d_f32. // Other variants were validated on the F16 im2col and keep it. bool pre_encode_f32_pointwise = false; + + // Run the pre_encode stack over overlapping time windows of this many + // output (subsampled) frames instead of the whole utterance; 0 = whole + // utterance. Exact (see build_pre_encode); bounds the stack's + // activations, which are otherwise linear in audio length (2 GB at + // 306 s for 256 channels). Ignored for masked or causal stacks. + int pre_encode_time_chunk = 0; }; // Resolve a conv-dispatch toggle from its env overrides. The DIRECT var forces @@ -146,6 +154,47 @@ bool resolve_conv_direct(const char * direct_env, const char * no_direct_env, bo // f32 weights. Metal and CPU prefer direct. bool detect_direct_pw(const char * backend); +// Default pre_encode time window, in output frames (80 ms each for the 8x +// stacks: ~20 s). Small enough that the stack's per-window activations +// (~150 MB at 256 channels) sit well under the encoder blocks' peak, large +// enough that the 1-frame overlap and the per-window kernel launches cost +// nothing measurable. +constexpr int kPreEncodeTimeChunk = 256; + +// Nodes a tiled pre_encode adds per window (conv stack + flatten + linear + +// concat); measured ~44, rounded up for headroom. Callers add +// kPreEncodeNodesPerWindow * pre_encode_window_count(...) to their graph size: +// the window count grows with audio length, so a fixed graph size overflows +// (GGML_ASSERT(cgraph->n_nodes < cgraph->size)) once the windows are many. +constexpr int kPreEncodeNodesPerWindow = 64; + +// Upper bound on the number of time windows. Windows exist to bound the +// stack's activations; past this many (>= 21 min of audio at the default +// chunk) making the windows longer costs nothing, while an unbounded window +// count would grow every graph/scheduler/metadata capacity with audio length. +constexpr int kPreEncodeMaxWindows = 64; + +// Graph nodes a tiled pre_encode can add at most (window cap * per-window): +// what a caller's graph and scheduler capacities need on top of their untiled +// size. +constexpr int kPreEncodeExtraGraphNodes = kPreEncodeNodesPerWindow * kPreEncodeMaxWindows; + +// Windows build_pre_encode builds for this mel length at this chunk (1 when it +// runs untiled). Same arithmetic as build_pre_encode's stage lengths. +int pre_encode_window_count(int64_t n_mel_frames, int chunk); + +// ggml metadata a tiled pre_encode needs for its windows, on top of an untiled +// one (measured 217 KB/window at 256 channels, rounded up). The windows also +// need their own tensors in the compute context, which does not grow by +// itself: without this the context pool runs out ("ggml_new_object: not +// enough space in the context's memory pool") at long audio or small chunks. +size_t pre_encode_meta_bytes(int64_t n_mel_frames, int chunk); + +// ConvPolicy::pre_encode_time_chunk with its env override: +// TRANSCRIBE_PRE_ENCODE_CHUNK=N sets the window (0 = whole utterance); +// otherwise `family_default`. +int resolve_pre_encode_time_chunk(int family_default); + // Scalar knobs for the block forward. All fields required except the // optional local-attention window (att_context_left / att_context_right). struct BlockParams { @@ -321,8 +370,9 @@ ggml_tensor * conv_1d_f32(ggml_context * ctx, int padding, int dilation); -// f32-friendly 2D depthwise conv. Use instead of ggml_conv_2d_dw / -// ggml_conv_2d_dw_direct for fp32 kernels on Metal (see conformer.cpp). +// f32-friendly 2D depthwise conv via im2col. Use instead of ggml_conv_2d_dw +// for fp32 kernels on Metal (see conformer.cpp); conv_2d_dw_direct_f32 is +// usually the better choice. ggml_tensor * conv_2d_dw_f32(ggml_context * ctx, ggml_tensor * kernel, ggml_tensor * data, @@ -348,7 +398,9 @@ ggml_tensor * conv_2d_dw_direct_f32(ggml_context * ctx, int d0, int d1); -// f32-friendly 1D depthwise conv. Same Metal reasoning. +// f32-friendly 1D depthwise conv. Same Metal reasoning. Switches algorithm +// on the batch size (im2col at B == 1, the direct op at B > 1), so callers +// with a batch-parity requirement use conv_2d_dw_direct_f32 instead. ggml_tensor * conv_1d_dw_f32(ggml_context * ctx, ggml_tensor * kernel, ggml_tensor * data, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 110645b0..09713315 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -366,6 +366,44 @@ transcribe_apply_warnings(transcribe_conv_pw_promote_unit) add_test(NAME transcribe_conv_pw_promote_unit COMMAND transcribe_conv_pw_promote_unit) +# ----------------------------------------------------------------------------- +# CONV_2D_DW (direct depthwise) on every device present +# ----------------------------------------------------------------------------- +# +# Evidence for the conformer depthwise policy: the direct op matches a host +# reference and the im2col helpers it replaces, is batch-invariant, and is +# supported on Metal and CPU (a regression fails rather than skips). + +add_executable(transcribe_conv_dw_direct_unit + conv_dw_direct_unit.cpp) + +target_link_libraries(transcribe_conv_dw_direct_unit PRIVATE transcribe ggml) + +target_include_directories(transcribe_conv_dw_direct_unit PRIVATE + ${CMAKE_SOURCE_DIR}/src) + +transcribe_apply_warnings(transcribe_conv_dw_direct_unit) + +add_test(NAME transcribe_conv_dw_direct_unit + COMMAND transcribe_conv_dw_direct_unit) + +# ----------------------------------------------------------------------------- +# build_pre_encode time tiling (ConvPolicy::pre_encode_time_chunk) is exact +# ----------------------------------------------------------------------------- + +add_executable(transcribe_pre_encode_tile_unit + pre_encode_tile_unit.cpp) + +target_link_libraries(transcribe_pre_encode_tile_unit PRIVATE transcribe ggml) + +target_include_directories(transcribe_pre_encode_tile_unit PRIVATE + ${CMAKE_SOURCE_DIR}/src) + +transcribe_apply_warnings(transcribe_pre_encode_tile_unit) + +add_test(NAME transcribe_pre_encode_tile_unit + COMMAND transcribe_pre_encode_tile_unit) + # ----------------------------------------------------------------------------- # causal_lm::pack_gate_up: alias / copy / mismatch, no model needed # ----------------------------------------------------------------------------- diff --git a/tests/conv_dw_direct_unit.cpp b/tests/conv_dw_direct_unit.cpp new file mode 100644 index 00000000..ec63b065 --- /dev/null +++ b/tests/conv_dw_direct_unit.cpp @@ -0,0 +1,334 @@ +// conv_dw_direct_unit.cpp - CONV_2D_DW (ggml_conv_2d_dw_direct) on every +// compute device present, against a host reference and the im2col forms. +// +// The conformer conv helpers used to route Metal to im2col + mul_mat on the +// strength of a comment saying CONV_2D_DW was unsupported there. The direct +// op is what keeps the depthwise convs from inflating the activation by the +// kernel width (the 1.9 GB granite block conv, the 2.3 GB conformer +// pre-encode conv at 306 s), so the policy rests on this test instead: +// +// * every device that reports CONV_2D_DW support computes the shapes the +// encoders use (k=15 1-D "same" conv, k=3 s=2 2-D pre-encode conv with +// the offline p=1 and the causal p=2 padding) within F32 rounding of a +// double-precision host reference, for F32 and F16 kernels; +// * Metal and CPU must report support: a regression there fails, it does +// not skip; +// * the direct op is batch-invariant: each utterance of an N=3 batch is +// bit-identical to the same utterance run alone (the batch-parity +// requirement conv_1d_dw_f32's B-dependent algorithm switch breaks); +// * the direct op agrees with the im2col helpers it replaces, within F16 +// rounding (the im2col routes round to F16 on some backends), where the +// device runs both. + +#include "conformer/conformer.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +namespace conf = transcribe::conformer; + +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) + +struct Case { + const char * name; + int kw, kh; // kernel + int w, h; // spatial input (W fast) + int c; // channels + int s0, s1, p0, p1; +}; + +// Shapes the encoders use, at toy sizes. Odd extents exercise the strided +// tail; the 1-D case is W=time, H=1 as granite / conformer blocks build it. +const Case kCases[] = { + { "block-1d-k15-same", 15, 1, 37, 1, 12, 1, 1, 7, 0 }, + { "pre-encode-k3s2-p1", 3, 3, 17, 23, 8, 2, 2, 1, 1 }, + { "pre-encode-k3s2-p2", 3, 3, 17, 23, 8, 2, 2, 2, 2 }, + { "pre-encode-k3s2-even", 3, 3, 16, 24, 8, 2, 2, 1, 1 }, +}; + +int64_t out_size(int64_t in, int k, int s, int p) { + return (in + 2 * p - k) / s + 1; +} + +uint32_t g_rng = 0x12345678u; + +float next_uniform() { + g_rng = g_rng * 1664525u + 1013904223u; + return static_cast(g_rng >> 8) / static_cast(1u << 24) * 2.0f - 1.0f; +} + +// Depthwise reference in double: out[n][c][oh][ow]. +std::vector reference(const Case & cs, + const std::vector & kernel, + const std::vector & data, + int n_batch) { + const int64_t ow = out_size(cs.w, cs.kw, cs.s0, cs.p0); + const int64_t oh = out_size(cs.h, cs.kh, cs.s1, cs.p1); + std::vector out(static_cast(ow * oh * cs.c * n_batch)); + for (int n = 0; n < n_batch; ++n) { + for (int c = 0; c < cs.c; ++c) { + for (int64_t y = 0; y < oh; ++y) { + for (int64_t x = 0; x < ow; ++x) { + double acc = 0.0; + for (int ky = 0; ky < cs.kh; ++ky) { + for (int kx = 0; kx < cs.kw; ++kx) { + const int64_t iy = y * cs.s1 - cs.p1 + ky; + const int64_t ix = x * cs.s0 - cs.p0 + kx; + if (iy < 0 || iy >= cs.h || ix < 0 || ix >= cs.w) { + continue; + } + const float d = data[static_cast(((n * cs.c + c) * cs.h + iy) * cs.w + ix)]; + const float k = kernel[static_cast((c * cs.kh + ky) * cs.kw + kx)]; + acc += static_cast(d) * static_cast(k); + } + } + out[static_cast(((n * cs.c + c) * oh + y) * ow + x)] = static_cast(acc); + } + } + } + } + return out; +} + +enum class Algo { Direct, Im2col }; + +struct RunResult { + bool supported = false; + std::vector out; +}; + +// Build and run one depthwise conv on `be`. Host kernel values must already +// be representable in `ktype` (F16 cases pre-round them). +RunResult run(ggml_backend_t be, + const Case & cs, + ggml_type ktype, + Algo algo, + const std::vector & kernel, + const std::vector & data, + int n_batch) { + RunResult res; + ggml_init_params ip = { 64 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true }; + ggml_context * ctx = ggml_init(ip); + if (ctx == nullptr) { + return res; + } + + ggml_tensor * k = ggml_new_tensor_4d(ctx, ktype, cs.kw, cs.kh, 1, cs.c); + ggml_tensor * d = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, cs.w, cs.h, cs.c, n_batch); + + ggml_tensor * out = nullptr; + if (algo == Algo::Direct) { + out = conf::conv_2d_dw_direct_f32(ctx, k, d, cs.s0, cs.s1, cs.p0, cs.p1, 1, 1); + } else if (cs.h == 1) { + // The 1-D helper takes [W, C, B]; it is single-utterance in its + // im2col form, so the caller only asks for n_batch == 1 here. + ggml_tensor * k3 = ggml_reshape_3d(ctx, k, cs.kw, 1, cs.c); + ggml_tensor * d3 = ggml_reshape_3d(ctx, d, cs.w, cs.c, n_batch); + out = conf::conv_1d_dw_f32(ctx, k3, d3, cs.s0, cs.p0, 1); + } else { + out = conf::conv_2d_dw_f32(ctx, k, d, cs.s0, cs.s1, cs.p0, cs.p1, 1, 1); + } + out = ggml_cont(ctx, out); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, out); + + // Every non-view node must run on this device, or the case is + // unsupported there (the caller decides whether that is a failure). + bool supported = true; + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + ggml_tensor * node = ggml_graph_node(gf, i); + if (node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_VIEW || + node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { + continue; + } + if (!ggml_backend_supports_op(be, node)) { + supported = false; + } + } + if (!supported) { + ggml_free(ctx); + return res; + } + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, be); + if (buf == nullptr) { + ggml_free(ctx); + return res; + } + + if (ktype == GGML_TYPE_F16) { + std::vector k16(kernel.size()); + for (size_t i = 0; i < kernel.size(); ++i) { + k16[i] = ggml_fp32_to_fp16(kernel[i]); + } + ggml_backend_tensor_set(k, k16.data(), 0, ggml_nbytes(k)); + } else { + ggml_backend_tensor_set(k, kernel.data(), 0, ggml_nbytes(k)); + } + ggml_backend_tensor_set(d, data.data(), 0, ggml_nbytes(d)); + + if (ggml_backend_graph_compute(be, gf) == GGML_STATUS_SUCCESS) { + res.supported = true; + res.out.resize(static_cast(ggml_nelements(out))); + ggml_backend_tensor_get(out, res.out.data(), 0, ggml_nbytes(out)); + } else { + std::fprintf(stderr, "FAIL graph compute on %s\n", ggml_backend_name(be)); + ++g_failures; + } + + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return res; +} + +// max |a-b| / (1 + |ref|) over all elements. +double max_rel(const std::vector & a, const std::vector & b) { + if (a.size() != b.size()) { + return INFINITY; + } + double worst = 0.0; + for (size_t i = 0; i < a.size(); ++i) { + const double e = std::fabs(static_cast(a[i]) - static_cast(b[i])) / + (1.0 + std::fabs(static_cast(b[i]))); + if (!(e <= worst)) { + worst = e; // also latches NaN + } + } + return worst; +} + +bool must_support(ggml_backend_dev_t dev) { + const enum ggml_backend_dev_type t = ggml_backend_dev_type(dev); + if (t == GGML_BACKEND_DEVICE_TYPE_CPU) { + return true; + } + const char * reg = ggml_backend_reg_name(ggml_backend_dev_backend_reg(dev)); + return reg != nullptr && std::strcmp(reg, "Metal") == 0; +} + +void run_device(ggml_backend_dev_t dev) { + const enum ggml_backend_dev_type t = ggml_backend_dev_type(dev); + if (t == GGML_BACKEND_DEVICE_TYPE_ACCEL) { + return; // BLAS etc.: not a standalone compute backend + } + ggml_backend_t be = ggml_backend_dev_init(dev, nullptr); + if (be == nullptr) { + std::fprintf(stderr, "note: could not init %s, skipped\n", ggml_backend_dev_name(dev)); + return; + } + const char * name = ggml_backend_name(be); + const bool required = must_support(dev); + + for (const Case & cs : kCases) { + for (ggml_type ktype : { GGML_TYPE_F32, GGML_TYPE_F16 }) { + const int n_batch = 3; + const int kn = cs.kw * cs.kh * cs.c; + const int dn = cs.w * cs.h * cs.c; + + std::vector kernel(static_cast(kn)); + for (float & v : kernel) { + v = next_uniform(); + if (ktype == GGML_TYPE_F16) { + v = ggml_fp16_to_fp32(ggml_fp32_to_fp16(v)); + } + } + std::vector data(static_cast(dn * n_batch)); + for (float & v : data) { + v = next_uniform(); + } + const std::vector ref = reference(cs, kernel, data, n_batch); + + const RunResult direct = run(be, cs, ktype, Algo::Direct, kernel, data, n_batch); + if (!direct.supported) { + if (required) { + std::fprintf(stderr, "FAIL %s: CONV_2D_DW unsupported for %s (%s kernel)\n", name, cs.name, + ggml_type_name(ktype)); + ++g_failures; + } else { + std::fprintf(stderr, "note: %s: CONV_2D_DW unsupported for %s (%s kernel), skipped\n", name, + cs.name, ggml_type_name(ktype)); + } + continue; + } + + // Direct vs the double reference: F32 rounding only. + const double e_ref = max_rel(direct.out, ref); + std::printf("%-8s %-22s %-4s direct vs ref %.2e\n", name, cs.name, ggml_type_name(ktype), e_ref); + CHECK(e_ref < 1e-5); + + // Batch invariance: utterance n of the batch == utterance n alone. + const size_t per_out = direct.out.size() / n_batch; + for (int n = 0; n < n_batch; ++n) { + std::vector one(data.begin() + static_cast(n) * dn, + data.begin() + static_cast(n + 1) * dn); + const RunResult solo = run(be, cs, ktype, Algo::Direct, kernel, one, 1); + CHECK(solo.supported); + CHECK(solo.out.size() == per_out); + if (solo.out.size() == per_out) { + const bool same = + std::memcmp(solo.out.data(), direct.out.data() + n * per_out, per_out * sizeof(float)) == 0; + if (!same) { + std::fprintf(stderr, "FAIL %s %s %s: batch item %d differs from solo run\n", name, cs.name, + ggml_type_name(ktype), n); + ++g_failures; + } + } + } + + // Direct vs the im2col helper it replaces (single utterance), a + // cross-check rather than the accuracy gate (that is the reference + // above): the F16-kernel im2col rounds the activations to F16, + // and Vulkan's im2col + mul_mat route carries F16-level error even + // for F32 kernels (~7e-4 against the reference), so the bound is + // the F16 one throughout. + std::vector first(data.begin(), data.begin() + dn); + const RunResult im2col = run(be, cs, ktype, Algo::Im2col, kernel, first, 1); + if (im2col.supported) { + std::vector direct0(direct.out.begin(), + direct.out.begin() + static_cast(per_out)); + const double e = max_rel(direct0, im2col.out); + const double tol = 5e-3; + std::printf("%-8s %-22s %-4s direct vs im2col %.2e\n", name, cs.name, ggml_type_name(ktype), e); + CHECK(e < tol); + } else { + std::fprintf(stderr, "note: %s: im2col form unsupported for %s (%s kernel)\n", name, cs.name, + ggml_type_name(ktype)); + } + } + } + ggml_backend_free(be); +} + +} // namespace + +int main() { + const size_t n_dev = ggml_backend_dev_count(); + for (size_t i = 0; i < n_dev; ++i) { + run_device(ggml_backend_dev_get(i)); + } + if (g_failures != 0) { + std::fprintf(stderr, "%d failure(s)\n", g_failures); + return EXIT_FAILURE; + } + std::printf("OK\n"); + return EXIT_SUCCESS; +} diff --git a/tests/pre_encode_tile_unit.cpp b/tests/pre_encode_tile_unit.cpp new file mode 100644 index 00000000..7a3b16b7 --- /dev/null +++ b/tests/pre_encode_tile_unit.cpp @@ -0,0 +1,229 @@ +// pre_encode_tile_unit.cpp - ConvPolicy::pre_encode_time_chunk is exact. +// +// build_pre_encode can run the DwStridingSubsampling stack over overlapping +// mel windows instead of the whole utterance (the stack's activations are +// otherwise linear in audio length: 3 GB at 306 s for cohere/canary-qwen). +// The window arithmetic (stride-8 aligned starts, one dropped frame on the +// left, balanced window lengths, the utterance edges left untouched) is the +// part that can silently go wrong, so this pins it on random weights: +// every window size, against the whole-utterance stack, across mel lengths +// that hit every residue mod 8 and the one-window / many-window boundaries, +// for batch 1 and 2, on every compute device present. + +#include "conformer/conformer.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml.h" + +#include +#include +#include +#include +#include + +namespace { + +namespace conf = transcribe::conformer; + +int g_failures = 0; + +constexpr int kMels = 16; // F' = 2 after three stride-2 stages +constexpr int kChan = 8; +constexpr int kModel = 12; + +uint32_t g_rng = 0x2468ace1u; + +float next_uniform() { + g_rng = g_rng * 1664525u + 1013904223u; + return static_cast(g_rng >> 8) / static_cast(1u << 24) * 2.0f - 1.0f; +} + +struct HostWeights { + std::vector conv0_w, conv0_b, conv2_w, conv2_b, conv3_w, conv3_b; + std::vector conv5_w, conv5_b, conv6_w, conv6_b, out_w, out_b; +}; + +std::vector rnd(size_t n, float scale) { + std::vector v(n); + for (float & x : v) { + x = next_uniform() * scale; + } + return v; +} + +// Runs build_pre_encode on `be` and returns the [d_model, T_enc, B] output, +// or an empty vector if the device cannot run the graph. +std::vector run(ggml_backend_t be, + const HostWeights & hw, + const std::vector & mel, + int t_mel, + int n_batch, + int chunk, + bool f32_paths, + int64_t * t_enc_out) { + ggml_init_params ip = { 4096 * ggml_tensor_overhead() + ggml_graph_overhead_custom(8192, false), nullptr, true }; + ggml_context * ctx = ggml_init(ip); + + auto new_t = [&](int64_t a, int64_t b, int64_t c, int64_t d) { + return ggml_new_tensor_4d(ctx, GGML_TYPE_F32, a, b, c, d); + }; + conf::PreEncodeView pe; + pe.conv0_w = new_t(3, 3, 1, kChan); + pe.conv0_b = new_t(kChan, 1, 1, 1); + pe.conv2_w = new_t(3, 3, 1, kChan); + pe.conv2_b = new_t(kChan, 1, 1, 1); + pe.conv3_w = new_t(1, 1, kChan, kChan); + pe.conv3_b = new_t(kChan, 1, 1, 1); + pe.conv5_w = new_t(3, 3, 1, kChan); + pe.conv5_b = new_t(kChan, 1, 1, 1); + pe.conv6_w = new_t(1, 1, kChan, kChan); + pe.conv6_b = new_t(kChan, 1, 1, 1); + pe.out_w = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, (kMels / 8) * kChan, kModel); + pe.out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kModel); + ggml_tensor * mel_in = new_t(t_mel, kMels, 1, n_batch); + ggml_tensor * weights[] = { pe.conv0_w, pe.conv0_b, pe.conv2_w, pe.conv2_b, pe.conv3_w, pe.conv3_b, + pe.conv5_w, pe.conv5_b, pe.conv6_w, pe.conv6_b, pe.out_w, pe.out_b }; + const std::vector * host[] = { &hw.conv0_w, &hw.conv0_b, &hw.conv2_w, &hw.conv2_b, &hw.conv3_w, &hw.conv3_b, + &hw.conv5_w, &hw.conv5_b, &hw.conv6_w, &hw.conv6_b, &hw.out_w, &hw.out_b }; + + conf::ConvPolicy policy{}; + policy.direct_dw_in_pre_encode = true; + policy.pre_encode_time_chunk = chunk; + // f32_paths: direct conv0 + F32-im2col pointwise convs. Otherwise the + // vendored ggml_conv_2d's F16 im2col (cohere / canary-qwen's config). + policy.direct_conv0_in_pre_encode = f32_paths; + policy.pre_encode_f32_pointwise = f32_paths; + + ggml_tensor * out = conf::build_pre_encode(ctx, pe, mel_in, policy, "pe", "pre_encode_tile_unit"); + std::vector res; + if (out == nullptr) { + std::fprintf(stderr, "FAIL build_pre_encode returned null (t_mel=%d chunk=%d)\n", t_mel, chunk); + ++g_failures; + ggml_free(ctx); + return res; + } + out = ggml_cont(ctx, out); + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 8192, false); + ggml_build_forward_expand(gf, out); + + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + if (!ggml_backend_supports_op(be, ggml_graph_node(gf, i))) { + ggml_free(ctx); + return res; // device cannot run it: caller skips + } + } + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, be); + for (size_t i = 0; i < sizeof(weights) / sizeof(weights[0]); ++i) { + ggml_backend_tensor_set(weights[i], host[i]->data(), 0, ggml_nbytes(weights[i])); + } + ggml_backend_tensor_set(mel_in, mel.data(), 0, ggml_nbytes(mel_in)); + if (ggml_backend_graph_compute(be, gf) == GGML_STATUS_SUCCESS) { + res.resize(static_cast(ggml_nelements(out))); + ggml_backend_tensor_get(out, res.data(), 0, ggml_nbytes(out)); + *t_enc_out = out->ne[1]; + } else { + std::fprintf(stderr, "FAIL graph compute (t_mel=%d chunk=%d)\n", t_mel, chunk); + ++g_failures; + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return res; +} + +void run_device(ggml_backend_dev_t dev, const HostWeights & hw) { + if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_ACCEL) { + return; + } + ggml_backend_t be = ggml_backend_dev_init(dev, nullptr); + if (be == nullptr) { + return; + } + const char * name = ggml_backend_name(be); + for (bool f32_paths : { false, true }) { + int n_cases = 0; + double worst = 0.0; + // Same inputs through the same per-element ops, so any difference is a + // backend choosing another matmul kernel for a different extent, i.e. + // accumulation order and, on Vulkan, the kernel's internal precision + // (up to ~2e-4 at the 1-2 frame windows here). Metal is bit-identical. + // A window-arithmetic error (wrong halo, misaligned start, dropped or + // duplicated frame) is O(0.1) on these weights, so the bound sits well + // between the two. + const double tol = 1e-3; + + // Mel lengths: every residue mod 8, around the window boundaries for the + // small windows below (T_enc = ceil(T_mel / 8)). + const int t_mels[] = { 8, 9, 31, 32, 33, 39, 40, 41, 47, 63, 64, 65, 100, 129, 203 }; + for (int n_batch : { 1, 2 }) { + for (int t_mel : t_mels) { + const std::vector mel = rnd(static_cast(t_mel) * kMels * n_batch, 2.0f); + int64_t t_ref = 0; + const std::vector ref = run(be, hw, mel, t_mel, n_batch, /*chunk=*/0, f32_paths, &t_ref); + if (ref.empty()) { + std::fprintf(stderr, "note: %s cannot run the pre_encode graph, skipped\n", name); + ggml_backend_free(be); + return; + } + for (int chunk : { 1, 2, 3, 4, 5, 7, 16 }) { + int64_t t_got = 0; + const std::vector got = run(be, hw, mel, t_mel, n_batch, chunk, f32_paths, &t_got); + ++n_cases; + if (t_got != t_ref || got.size() != ref.size()) { + std::fprintf(stderr, "FAIL %s t_mel=%d B=%d chunk=%d: T_enc %lld vs %lld\n", name, t_mel, + n_batch, chunk, static_cast(t_got), static_cast(t_ref)); + ++g_failures; + continue; + } + double e = 0.0; + for (size_t i = 0; i < ref.size(); ++i) { + const double d = std::fabs(static_cast(got[i]) - static_cast(ref[i])) / + (1.0 + std::fabs(static_cast(ref[i]))); + if (!(d <= e)) { + e = d; + } + } + if (!(e <= worst)) { + worst = e; + } + if (!(e < tol)) { + std::fprintf(stderr, "FAIL %s t_mel=%d B=%d chunk=%d f32=%d: max rel %.3e\n", name, t_mel, + n_batch, chunk, f32_paths ? 1 : 0, e); + ++g_failures; + } + } + } + } + std::printf("%-8s %-9s %d tiled cases, worst rel diff vs whole-utterance %.2e\n", name, + f32_paths ? "f32" : "f16-im2col", n_cases, worst); + } + ggml_backend_free(be); +} + +} // namespace + +int main() { + HostWeights hw; + hw.conv0_w = rnd(3 * 3 * kChan, 0.5f); + hw.conv0_b = rnd(kChan, 0.2f); + hw.conv2_w = rnd(3 * 3 * kChan, 0.5f); + hw.conv2_b = rnd(kChan, 0.2f); + hw.conv3_w = rnd(kChan * kChan, 0.5f); + hw.conv3_b = rnd(kChan, 0.2f); + hw.conv5_w = rnd(3 * 3 * kChan, 0.5f); + hw.conv5_b = rnd(kChan, 0.2f); + hw.conv6_w = rnd(kChan * kChan, 0.5f); + hw.conv6_b = rnd(kChan, 0.2f); + hw.out_w = rnd(static_cast(kMels / 8) * kChan * kModel, 0.5f); + hw.out_b = rnd(kModel, 0.2f); + + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + run_device(ggml_backend_dev_get(i), hw); + } + if (g_failures != 0) { + std::fprintf(stderr, "%d failure(s)\n", g_failures); + return EXIT_FAILURE; + } + std::printf("OK\n"); + return EXIT_SUCCESS; +}