From e16cac3c24e9e0a52a538bef5bb1c97c856e7075 Mon Sep 17 00:00:00 2001 From: Ali Hamdi Ali Fadel Date: Sat, 24 Oct 2026 10:29:03 +0000 Subject: [PATCH] conformer: make the GLU sigmoid input contiguous so CUDA keeps it on the GPU --- src/arch/granite/encoder.cpp | 9 +-- src/arch/granite5_ctc/encoder.cpp | 11 +--- src/arch/granite_nar/encoder.cpp | 6 +- src/conformer/conformer.cpp | 24 ++++---- src/conformer/conformer.h | 4 ++ tests/CMakeLists.txt | 17 ++++++ tests/conformer_glu_unit.cpp | 95 +++++++++++++++++++++++++++++++ 7 files changed, 133 insertions(+), 33 deletions(-) create mode 100644 tests/conformer_glu_unit.cpp diff --git a/src/arch/granite/encoder.cpp b/src/arch/granite/encoder.cpp index 1396df66a..a21ac5847 100644 --- a/src/arch/granite/encoder.cpp +++ b/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)); diff --git a/src/arch/granite5_ctc/encoder.cpp b/src/arch/granite5_ctc/encoder.cpp index 79d1efe64..78721129c 100644 --- a/src/arch/granite5_ctc/encoder.cpp +++ b/src/arch/granite5_ctc/encoder.cpp @@ -262,7 +262,6 @@ 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); @@ -270,14 +269,8 @@ ggml_tensor * conv_module(ggml_context * ctx, 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 diff --git a/src/arch/granite_nar/encoder.cpp b/src/arch/granite_nar/encoder.cpp index ea7025f5d..1529c2c9c 100644 --- a/src/arch/granite_nar/encoder.cpp +++ b/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 diff --git a/src/conformer/conformer.cpp b/src/conformer/conformer.cpp index 9b6c2ded0..0c1547714 100644 --- a/src/conformer/conformer.cpp +++ b/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)); diff --git a/src/conformer/conformer.h b/src/conformer/conformer.h index 745416343..32141ec38 100644 --- a/src/conformer/conformer.h +++ b/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. diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 110645b02..35be23967 100644 --- a/tests/CMakeLists.txt +++ b/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 # ----------------------------------------------------------------------------- diff --git a/tests/conformer_glu_unit.cpp b/tests/conformer_glu_unit.cpp new file mode 100644 index 000000000..4df065cfa --- /dev/null +++ b/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 +#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) + +} // 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 input(2 * C * T * B); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = 0.25f * static_cast(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 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; +}