conformer: make the GLU sigmoid input contiguous so CUDA keeps it on the GPU

src/arch/granite/encoder.cpp

@@ -225,7 +225,6 @@ ggml_tensor * granite_conv_module(ggml_context *          ctx,
                                   int                     conv_kernel,
                                   int                     inner_dim) {
     const int64_t d_model = x->ne[0];
-    const int64_t T       = x->ne[1];
 
     // Pre-LayerNorm with bias.
     x = layer_norm(ctx, x, b.norm_conv_w, b.norm_conv_b);
@@ -239,13 +238,7 @@ ggml_tensor * granite_conv_module(ggml_context *          ctx,
         x                 = ggml_add(ctx, x, b.conv_pointwise1_b);
     }
 
-    // GLU: split ne[0] in half, gate * sigmoid(value).
-    {
-        ggml_tensor * gate  = ggml_view_2d(ctx, x, inner_dim, T, x->nb[1], /*offset=*/0);
-        ggml_tensor * value = ggml_view_2d(ctx, x, inner_dim, T, x->nb[1], inner_dim * ggml_element_size(x));
-        x                   = ggml_mul(ctx, gate, ggml_sigmoid(ctx, value));
-    }
-    // x ne = [inner_dim, T]
+    x = transcribe::conformer::glu(ctx, x);  // [inner_dim, T]
 
     // Transpose for depthwise conv: [inner_dim, T] -> [T, inner_dim].
     x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));

src/arch/granite5_ctc/encoder.cpp

@@ -262,22 +262,15 @@ ggml_tensor * conv_module(ggml_context *              ctx,
                           int                         stride,
                           ggml_tensor *               frame_mask) {
     const int64_t d_model = x->ne[0];
-    const int64_t T       = x->ne[1];
     const int64_t B       = x->ne[2];
 
     x = layer_norm(ctx, x, b.norm_conv_w, b.norm_conv_b);
 
     x = ggml_mul_mat(ctx, b.conv_pointwise1_w, x);  // [2*inner, T, B]
     x = ggml_add(ctx, x, b.conv_pointwise1_b);
 
-    // GLU over ne[0]: first half * sigmoid(second half). Matches
-    // F.glu(dim=-1) on the reference's [B, T, 2*inner] layout.
-    {
-        ggml_tensor * gate = ggml_view_3d(ctx, x, inner_dim, T, B, x->nb[1], x->nb[2], /*offset=*/0);
-        ggml_tensor * value =
-            ggml_view_3d(ctx, x, inner_dim, T, B, x->nb[1], x->nb[2], inner_dim * ggml_element_size(x));
-        x = ggml_mul(ctx, gate, ggml_sigmoid(ctx, value));
-    }
+    // Matches F.glu(dim=-1) on the reference's [B, T, 2*inner] layout.
+    x = transcribe::conformer::glu(ctx, x);
 
     // Zero the padded frames before the depthwise conv, matching the
     // reference's `masked_fill(~attention_mask.unsqueeze(-1), 0.0)`. This

src/arch/granite_nar/encoder.cpp

@@ -183,11 +183,7 @@ ggml_tensor * conv_module(ggml_context *             ctx,
         x                 = ggml_mul_mat(ctx, pw1, x);
         x                 = ggml_add(ctx, x, b.conv_pointwise1_b);
     }
-    {
-        ggml_tensor * gate  = ggml_view_2d(ctx, x, inner_dim, T, x->nb[1], 0);
-        ggml_tensor * value = ggml_view_2d(ctx, x, inner_dim, T, x->nb[1], inner_dim * ggml_element_size(x));
-        x                   = ggml_mul(ctx, gate, ggml_sigmoid(ctx, value));
-    }
+    x                 = transcribe::conformer::glu(ctx, x);
     x                 = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
     const int padding = (conv_kernel - 1) / 2;
     // Direct depthwise op, not conv_1d_dw_f32's im2col + mul_mat: im2col

src/conformer/conformer.cpp

@@ -78,6 +78,18 @@ ggml_tensor * macaron_ff_residual(ggml_context * ctx,
     return ggml_add(ctx, x, y);
 }
 
+// GLU over ne[0]. Both halves are strided views of x. ggml-cuda runs unary
+// ops only on contiguous sources, so sigmoid on the strided half would be
+// scheduled on the CPU, with a device round trip in every conv module.
+ggml_tensor * glu(ggml_context * ctx, ggml_tensor * x) {
+    assert(x->ne[0] % 2 == 0);
+    const int64_t half  = x->ne[0] / 2;
+    ggml_tensor * gate  = ggml_view_4d(ctx, x, half, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], 0);
+    ggml_tensor * value = ggml_view_4d(ctx, x, half, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3],
+                                       half * ggml_element_size(x));
+    return ggml_mul(ctx, gate, ggml_sigmoid(ctx, ggml_cont(ctx, value)));
+}
+
 // Shaw / Transformer-XL relative-position skew. Rotates each row of the
 // [pos_len, T_q, H, 1] score matrix so column k holds relative offset k.
 ggml_tensor * rel_shift(ggml_context * ctx, ggml_tensor * x) {
@@ -385,17 +397,7 @@ ggml_tensor * conv_module(ggml_context * ctx, ggml_tensor * x, const BlockView &
             }
         }
 
-        // GLU: split ne[0] in half, gate * sigmoid(value). The views carry
-        // the batch axis (ne[2]) so the split is per-utterance.
-        {
-            const int64_t T     = x->ne[1];
-            const int64_t half  = x->ne[0] / 2;
-            ggml_tensor * gate  = ggml_view_3d(ctx, x, half, T, B, x->nb[1], x->nb[2],
-                                               /*offset=*/0);
-            ggml_tensor * value = ggml_view_3d(ctx, x, half, T, B, x->nb[1], x->nb[2], half * ggml_element_size(x));
-            x                   = ggml_mul(ctx, gate, ggml_sigmoid(ctx, value));
-        }
-        // x ne = [d_model, T, B]
+        x = glu(ctx, x);  // [d_model, T, B]
 
         // Transpose for depthwise conv: [d_model, T, B] -> [T, d_model, B].
         x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));

src/conformer/conformer.h

@@ -302,6 +302,10 @@ ggml_tensor * macaron_ff_residual(ggml_context * ctx,
                                   ggml_tensor *  lin2_w,
                                   ggml_tensor *  lin2_b);
 
+// GLU over ne[0]: first half * sigmoid(second half), as F.glu on the
+// channel axis. [2*C, T, B] -> [C, T, B].
+ggml_tensor * glu(ggml_context * ctx, ggml_tensor * x);
+
 // Shaw / Transformer-XL relative-position skew. Turns a
 // [pos_len, T_q, H, 1] score matrix into a [pos_len, T_q, H, 1]
 // matrix rotated so column k holds the score for relative offset k.

tests/CMakeLists.txt

@@ -366,6 +366,23 @@ transcribe_apply_warnings(transcribe_conv_pw_promote_unit)
 add_test(NAME transcribe_conv_pw_promote_unit
     COMMAND transcribe_conv_pw_promote_unit)
 
+# -----------------------------------------------------------------------------
+# Conformer GLU helper (pure graph, no model)
+# -----------------------------------------------------------------------------
+
+add_executable(transcribe_conformer_glu_unit
+    conformer_glu_unit.cpp)
+
+target_link_libraries(transcribe_conformer_glu_unit PRIVATE transcribe ggml)
+
+target_include_directories(transcribe_conformer_glu_unit PRIVATE
+    ${CMAKE_SOURCE_DIR}/src)
+
+transcribe_apply_warnings(transcribe_conformer_glu_unit)
+
+add_test(NAME transcribe_conformer_glu_unit
+    COMMAND transcribe_conformer_glu_unit)
+
 # -----------------------------------------------------------------------------
 # causal_lm::pack_gate_up: alias / copy / mismatch, no model needed
 # -----------------------------------------------------------------------------

tests/conformer_glu_unit.cpp

@@ -0,0 +1,95 @@
+// conformer_glu_unit.cpp - unit test for transcribe::conformer::glu.
+//
+// Checks the GLU values on a batched [2*C, T, B] input, and that the sigmoid
+// reads a contiguous tensor: ggml-cuda runs unary ops only on contiguous
+// sources, so a strided input would move the sigmoid to the CPU.
+
+#include "conformer/conformer.h"
+#include "ggml-backend.h"
+#include "ggml.h"
+
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+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)
+
+}  // namespace
+
+int main() {
+    ggml_backend_t backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
+    if (backend == nullptr) {
+        std::fprintf(stderr, "SKIP: could not initialize CPU backend\n");
+        return 77;
+    }
+
+    ggml_init_params params{};
+    params.mem_size    = 1024 * 1024;
+    params.no_alloc    = true;
+    ggml_context * ctx = ggml_init(params);
+    if (ctx == nullptr) {
+        ggml_backend_free(backend);
+        return EXIT_FAILURE;
+    }
+
+    const int     C = 4;
+    const int     T = 3;
+    const int     B = 2;
+    ggml_tensor * x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2 * C, T, B);
+    ggml_tensor * y = transcribe::conformer::glu(ctx, x);
+
+    CHECK(y->ne[0] == C && y->ne[1] == T && y->ne[2] == B);
+    const ggml_tensor * sigmoid = y->src[1];
+    CHECK(sigmoid->op == GGML_OP_UNARY && ggml_get_unary_op(sigmoid) == GGML_UNARY_OP_SIGMOID);
+    CHECK(ggml_is_contiguous(sigmoid->src[0]));
+
+    ggml_cgraph * graph = ggml_new_graph(ctx);
+    ggml_build_forward_expand(graph, y);
+
+    ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend);
+    if (buffer == nullptr) {
+        ggml_free(ctx);
+        ggml_backend_free(backend);
+        return EXIT_FAILURE;
+    }
+
+    std::vector<float> input(2 * C * T * B);
+    for (size_t i = 0; i < input.size(); ++i) {
+        input[i] = 0.25f * static_cast<float>(i) - 3.0f;
+    }
+    ggml_backend_tensor_set(x, input.data(), 0, input.size() * sizeof(float));
+
+    CHECK(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS);
+    std::vector<float> output(C * T * B);
+    ggml_backend_tensor_get(y, output.data(), 0, output.size() * sizeof(float));
+    for (int row = 0; row < T * B; ++row) {
+        for (int c = 0; c < C; ++c) {
+            const float gate     = input[row * 2 * C + c];
+            const float value    = input[row * 2 * C + C + c];
+            const float expected = gate / (1.0f + std::exp(-value));
+            CHECK(std::fabs(output[row * C + c] - expected) <= 1e-5f * (1.0f + std::fabs(expected)));
+        }
+    }
+
+    ggml_backend_buffer_free(buffer);
+    ggml_free(ctx);
+    ggml_backend_free(backend);
+
+    if (g_failures != 0) {
+        std::fprintf(stderr, "conformer_glu_unit: %d failures\n", g_failures);
+        return EXIT_FAILURE;
+    }
+    std::fprintf(stdout, "conformer_glu_unit: ok\n");
+    return EXIT_SUCCESS;
+}