conformer: direct depthwise conv everywhere

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<size_t>(conf::kPreEncodeNodesPerWindow) *
+                    static_cast<size_t>(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;

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<size_t>(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<int>(cm->plan.scheduler_list.size()), 16384, /*parallel=*/false,
+                                           static_cast<int>(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<int>(cm->plan.scheduler_list.size()), 16384, false, true);
+                                           static_cast<int>(cm->plan.scheduler_list.size()),
+                                           16384 + conformer::kPreEncodeExtraGraphNodes, false, true);
         if (cc->sched == nullptr) {
             return TRANSCRIBE_ERR_BACKEND;
         }

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<size_t>(conf::kPreEncodeNodesPerWindow) *
+                    static_cast<size_t>(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;

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<int>(cm->plan.scheduler_list.size()), 16384, false, true);
+                                           static_cast<int>(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<int>(cm->plan.scheduler_list.size()), 16384, false, true);
+                                           static_cast<int>(cm->plan.scheduler_list.size()),
+                                           16384 + conformer::kPreEncodeExtraGraphNodes, false, true);
         if (cc->sched == nullptr) {
             return TRANSCRIBE_ERR_BACKEND;
         }

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.

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 ------------------------------------------------------

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;
 }