1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
// arch/voxtral_realtime/encoder.h - causal RoPE audio encoder + projector.
//
// VoxtralRealtimeEncoder: a 2-layer LEFT-PAD causal Conv1d stem (GELU) then
// 32 pre-norm RMSNorm transformer blocks with NEOX RoPE (theta 1e6, head_dim
// 64), causal + sliding-window(750) attention (q/v/out bias, k none) and a
// SwiGLU/silu MLP (bias only on down). A final RMSNorm yields enc.out; the
// projector groups 4 consecutive frames and runs Linear -> GELU -> Linear.
//
// One graph processes the WHOLE clip offline: mel [n_mels, n_mel_frames] in,
// proj.out [dec_hidden, n_audio] out. positions + a host-prepared
// sliding-window-causal mask are graph inputs.
#pragma once
#include "causal_lm/causal_lm.h"
#include "ggml.h"
#include <vector>
struct ggml_context;
struct ggml_tensor;
struct ggml_cgraph;
namespace transcribe::voxtral_realtime {
struct HParams;
struct Weights;
struct EncoderDumps {
ggml_tensor * embedder_out = nullptr; // conv stem output [d_model, T_enc]
ggml_tensor * enc_out = nullptr; // final RMSNorm output
ggml_tensor * proj_out = nullptr; // projector output (audio embeds)
std::vector<ggml_tensor *> block_outs; // per-block residual outputs
};
struct EncoderBuild {
ggml_tensor * mel_in = nullptr; // [n_mels, n_mel_frames] input
ggml_tensor * positions_in = nullptr; // [T_enc] i32 RoPE positions
ggml_tensor * mask_in = nullptr; // [T_enc, T_enc] f16 sw-causal mask
ggml_tensor * out = nullptr; // proj.out [dec_hidden, n_audio]
EncoderDumps dumps {};
ggml_cgraph * graph = nullptr;
int T_enc = 0;
int n_audio = 0;
};
// Build the encoder+projector forward graph for `n_mel_frames` mel frames.
// The stride-2 causal conv2 yields T_enc = conv-out frames; n_audio = T_enc/4.
EncoderBuild build_encoder_graph(ggml_context * ctx,
const Weights & weights,
const HParams & hp,
int n_mel_frames,
bool use_flash);
// ---------------------------------------------------------------------------
// Batched offline encoder (transcribe_run_batch fast path)
// ---------------------------------------------------------------------------
//
// Mirrors the reference, which runs the audio tower batched over [B, T, C] with
// NO audio attention_mask: every clip's mel is RIGHT-padded to the batch-max
// `n_mel_frames` and stacked on ne[2]=B. Because the conv stem is causal
// (left-pad) and attention is causal + sliding-window, a real frame at position
// t only attends to <= t, so the right-padding of shorter clips never
// contaminates their real [0, T_enc_b) outputs. A SINGLE `positions [T_enc]`
// and sw-causal `mask [T_enc, T_enc]` therefore serve every row; the caller
// slices each row's real n_audio_b out of out[:, :, b]. `use_flash` selects the
// flash vs manual-softmax attention (honor the session's encoder_use_flash, so
// the default CPU source-of-truth non-flash path batches too).
struct EncoderBuildBatched {
ggml_tensor * mel_in = nullptr; // [n_mels, n_mel_frames, B] input
ggml_tensor * positions_in = nullptr; // [T_enc] i32 RoPE positions (shared)
ggml_tensor * mask_in = nullptr; // [T_enc, T_enc] f16 sw-causal (shared)
ggml_tensor * out = nullptr; // proj.out [dec_hidden, n_audio, B]
ggml_cgraph * graph = nullptr;
int T_enc = 0;
int n_audio = 0;
int n_batch = 0;
};
EncoderBuildBatched build_encoder_graph_batched(ggml_context * ctx,
const Weights & weights,
const HParams & hp,
int n_mel_frames,
int n_batch,
bool use_flash);
// ---------------------------------------------------------------------------
// Incremental streaming encoder (matches the reference StaticCache mechanism)
// ---------------------------------------------------------------------------
//
// The reference runs the encoder transformer INCREMENTALLY: each decode step
// feeds `downsample_factor` new embedder frames through the 32 transformer
// layers against an encoder StaticCache (sliding_window=750), producing one
// audio token. We split that into two graphs:
//
// build_embedder_graph — conv stem only: mel -> embedder frames
// [d_model, T_enc]. Causal convs (left-pad), so the
// frames are append-only/stable (a whole-buffer
// recompute equals the reference per-chunk conv +
// padding cache, frame-for-frame).
// build_encoder_chunk_graph — the expensive part, INCREMENTAL: `n_new`
// embedder frames as queries attend to the encoder
// KV cache (prior frames' K/V) under a sliding-
// window-causal mask; their K/V are appended at
// cache offset `n_kv_before`. Final RMSNorm +
// projector group-of-4 -> audio embeds.
struct EmbedderBuild {
ggml_tensor * mel_in = nullptr; // [n_mels, n_mel_frames] input
ggml_tensor * out = nullptr; // [d_model, T_enc] embedder frames
ggml_cgraph * graph = nullptr;
int T_enc = 0;
};
// Conv stem only (no transformer). T_enc = (n_mel_frames - 2)/2 + 1.
EmbedderBuild build_embedder_graph(ggml_context * ctx,
const Weights & weights,
const HParams & hp,
int n_mel_frames);
// Conv stem with the streaming PADDING CACHE (reference VoxtralRealtimeConv1dCacheLayer).
// Feeds only `n_new_mel` NEW mel frames; the conv left-context comes from the host-
// carried caches instead of zero left-pad: conv1 (k3 s1, left_pad 2) ← last 2 mel
// frames; conv2 (k3 s2, left_pad 1) ← last 1 conv1-output frame. Both zero on the
// first chunk, which makes the cached path bit-identical to the whole-buffer left-pad
// conv. Emits M_emb = n_new_mel/2 new embedder frames (n_new_mel must be even) plus
// `cache2_out` = the chunk's last conv1-output frame for the next call. The conv1
// cache (last 2 mel frames) is trivially the input tail, carried host-side.
struct EmbedderChunkBuild {
ggml_tensor * mel_in = nullptr; // [n_mels, n_new_mel] new mel frames
ggml_tensor * cache1_in = nullptr; // [n_mels, 2] prev 2 mel frames (conv1 left ctx)
ggml_tensor * cache2_in = nullptr; // [d_model, 1] prev last conv1-out frame (conv2 left ctx)
ggml_tensor * out = nullptr; // [d_model, M_emb] new embedder frames
ggml_tensor * cache2_out = nullptr; // [d_model, 1] new last conv1-out frame (next cache2)
ggml_cgraph * graph = nullptr;
int M_emb = 0;
};
EmbedderChunkBuild build_embedder_chunk_graph(ggml_context * ctx,
const Weights & weights,
const HParams & hp,
int n_new_mel);
struct EncoderChunkBuild {
ggml_tensor * embed_in = nullptr; // [d_model, n_new] new embedder frames
ggml_tensor * positions_in = nullptr; // [n_new] i32 ABSOLUTE enc-frame pos
ggml_tensor * mask_in = nullptr; // [read_len, n_new] f16 sw-causal
ggml_tensor * out = nullptr; // proj.out audio embeds [dec_hidden, n_audio_new]
ggml_tensor * enc_out = nullptr; // [d_model, n_new] final-normed enc frames
ggml_cgraph * graph = nullptr;
int n_new = 0;
int read_len = 0;
int n_audio_new = 0;
};
// Build one incremental encoder chunk against the encoder KV cache ring. Writes
// the `n_new` frames' K/V into cache rows [write_slot, write_slot + n_new), then
// reads the `read_len`-row window starting at `read_start` for attention (the
// host mask enforces the 750-frame sliding window using ABSOLUTE positions). The
// caller owns the ring (compaction + slot bookkeeping). `n_new` must be a
// multiple of proj_downsample (4) so projector groups never straddle chunks.
EncoderChunkBuild build_encoder_chunk_graph(ggml_context * ctx,
const Weights & weights,
const HParams & hp,
transcribe::causal_lm::KvCache & enc_kv,
int n_new,
int write_slot,
int read_start,
int read_len,
bool use_flash);
} // namespace transcribe::voxtral_realtime