diff --git a/CMakeLists.txt b/CMakeLists.txt index 94a395c8..0b4896fd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -875,6 +875,14 @@ if (ENGINE_BUILD_TESTS) COMMAND audio_chunking_test ) + add_engine_unittest(voxtral_realtime_stream_chunking_test tests/unittests/test_voxtral_realtime_stream_chunking.cpp) + target_include_directories(voxtral_realtime_stream_chunking_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/unittests) + + add_test( + NAME voxtral_realtime_stream_chunking_test + COMMAND voxtral_realtime_stream_chunking_test + ) + add_engine_unittest(chinese_normalization_test tests/unittests/test_chinese_normalization.cpp) add_test( diff --git a/docs/asr.md b/docs/asr.md index d681b3b2..f31bedaa 100644 --- a/docs/asr.md +++ b/docs/asr.md @@ -286,6 +286,31 @@ Streaming CLI: audiocpp_cli --task asr --family voxtral_realtime --model models/Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf --backend cuda --threads 8 --mode streaming --audio assets/resources/sample.wav --text-out transcript.txt ``` +> **Throughput.** A streaming step always advances 80 ms of audio, so a step has to cost under +> 80 ms to keep up with a realtime source. Measured on an Apple M3 Air (Metal, q8_0): +> +> | Config | short clip, cool | sustained 7 min | +> |---|---:|---:| +> | default | 78 ms/step (0.98x) | 88 ms/step (1.10x) | +> | `stream_batch_tokens=4` | 60 ms/step (0.76x) | 74 ms/step (0.92x) | +> +> The default splits roughly 48 ms for the text decoder and 30 ms for the audio encoder; batching +> takes the encoder to ~13 ms. The second column is what a long session actually gets on a fanless +> machine: a short clip run immediately after the 7-minute one still measured 88 ms/step, so the +> gap is the machine staying warm rather than anything that resets between sessions. Budget for the +> sustained column, and prefer `stream_batch_tokens=4` if the source is realtime. +> +> The decoder runs one step per 80 ms whether the audio holds speech or silence, so a session that +> does fall behind stays behind — the lag is monotonic and does not recover during pauses. Measure +> your own hardware before relying on a live source. + +Streaming session options: + +| Option | Default | Meaning | +|---|---:|---| +| `--session-option voxtral_realtime.stream_batch_tokens=` | `1` | Audio tokens per encoder forward. The decoder still runs one step per 80 ms; batching only amortizes the encoder's fixed per-forward cost, which dominates it. `4` takes the encoder from ~30 to ~13 ms/step, at the price of delaying every partial by up to `n * 80 ms`. | +| `--session-option voxtral_realtime.stream_decode_cache_steps=` | `1024` | Decoder KV cache size in 80 ms steps (~82 s of context). Built once when the stream starts, so a long session never stalls on a cache-growth rebuild; the cache ring then wraps in place, and a 7-minute stream stays coherent across five wraparounds. Lower values trade context for memory, not for speed. | + Streaming server config: ```json @@ -328,11 +353,16 @@ curl -N http://127.0.0.1:8080/v1/audio/transcriptions \ | `--top-k` | integer | `50` | Top-k sampling limit; `0` disables top-k. | | `--seed` | integer | `1234` | Sampling seed. | | `--text-out` | TXT path | not set | Transcript output. The transcript is also printed to stdout. | -| `--session-option voxtral_realtime.weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | Shared matmul weight storage type. | -| `--session-option voxtral_realtime.audio_encoder_weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0` | shared setting | Audio encoder matmul weight storage type. | -| `--session-option voxtral_realtime.text_decoder_weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0` | shared setting | Text decoder matmul weight storage type. | +| `--session-option voxtral_realtime.weight_type=` | `native`, `f32`, `f16`, `bf16`, `q4_0`, `q4_k`, `q5_k`, `q6_k`, `q8_0` | `native` | Shared matmul weight storage type. | +| `--session-option voxtral_realtime.audio_encoder_weight_type=` | same as above | shared setting | Audio encoder matmul weight storage type. Leave at `native` for streaming: the encoder is not bandwidth-bound there, so quantizing it makes it slower. | +| `--session-option voxtral_realtime.text_decoder_weight_type=` | same as above | shared setting | Text decoder matmul weight storage type. `q4_k` roughly halves the streaming decoder step cost. | | `--session-option voxtral_realtime.audio_encoder_graph_arena_mb=` | MB | `512` | Audio encoder graph arena size. | | `--session-option voxtral_realtime.text_decoder_prefill_graph_arena_mb=` | MB | `512` | Text decoder prefill graph arena size. | | `--session-option voxtral_realtime.text_decoder_decode_graph_arena_mb=` | MB | `512` | Text decoder cached-step graph arena size. | +Weight storage types are applied when the model loads, so asking for one the GGUF does not already +hold means requantizing on the CPU before the first token appears — around three minutes for +`q4_k` from the shipped q8_0 package. Prefer the published `q4_k` GGUF variant, which needs no +load-time conversion. See [GGUF](gguf.md). + For backend weight-type controls, use `audiocpp_cli --inspect --model --family `. diff --git a/include/engine/models/voxtral_realtime/audio_encoder.h b/include/engine/models/voxtral_realtime/audio_encoder.h index dbc6e5b7..bd161e81 100644 --- a/include/engine/models/voxtral_realtime/audio_encoder.h +++ b/include/engine/models/voxtral_realtime/audio_encoder.h @@ -27,6 +27,8 @@ class VoxtralRealtimeAudioEncoderRuntime { VoxtralRealtimeAudioEmbeddings encode_stream_chunk( const VoxtralRealtimeFeatures & features, VoxtralRealtimeAudioEncoderStreamState & state); + // Emits the summed upload/compute/readback split of every streaming chunk seen so far. + void log_stream_timings() const; private: std::unique_ptr impl_; diff --git a/include/engine/models/voxtral_realtime/frontend.h b/include/engine/models/voxtral_realtime/frontend.h index 23da2979..4827d7c1 100644 --- a/include/engine/models/voxtral_realtime/frontend.h +++ b/include/engine/models/voxtral_realtime/frontend.h @@ -21,15 +21,18 @@ class VoxtralRealtimeFrontend { explicit VoxtralRealtimeFrontend(std::shared_ptr assets); VoxtralRealtimeFeatures extract(const runtime::AudioBuffer & audio, bool first_chunk) const; + // A steady chunk carries `steady_tokens` audio tokens of 80 ms each; batching them into one + // extraction yields 8 * steady_tokens feature frames from a single STFT pass. VoxtralRealtimeFeatures extract_stream_chunk( const runtime::AudioBuffer & audio, bool first_chunk, - VoxtralRealtimeFrontendStreamState & state) const; + VoxtralRealtimeFrontendStreamState & state, + int64_t steady_tokens = 1) const; int64_t first_stream_chunk_samples() const; - int64_t steady_stream_chunk_samples() const; + int64_t steady_stream_chunk_samples(int64_t steady_tokens = 1) const; int64_t first_stream_chunk_advance_samples() const; - int64_t steady_stream_chunk_advance_samples() const; + int64_t steady_stream_chunk_advance_samples(int64_t steady_tokens = 1) const; private: std::shared_ptr assets_; diff --git a/include/engine/models/voxtral_realtime/session.h b/include/engine/models/voxtral_realtime/session.h index af4f8d3e..06e7bb7f 100644 --- a/include/engine/models/voxtral_realtime/session.h +++ b/include/engine/models/voxtral_realtime/session.h @@ -44,6 +44,9 @@ class VoxtralRealtimeSession final runtime::TaskResult run_single(const VoxtralRealtimeRequest & request, bool first_chunk); runtime::StreamEvent process_available_stream_chunks(); runtime::StreamEvent process_one_stream_chunk(const runtime::AudioBuffer & audio); + // Feeds a freshly decoded token back as the next stream input and, when it carries text, + // appends it to the transcript and refreshes the event's partial. + void record_stream_token(int32_t token, runtime::StreamEvent & event); runtime::TaskSpec task_; std::shared_ptr assets_; @@ -54,6 +57,8 @@ class VoxtralRealtimeSession final size_t text_decoder_weight_context_bytes_ = 128ull * 1024ull * 1024ull; assets::TensorStorageType audio_encoder_weight_storage_type_ = assets::TensorStorageType::Native; assets::TensorStorageType text_decoder_weight_storage_type_ = assets::TensorStorageType::Native; + int64_t stream_decode_cache_steps_ = 1024; + int64_t stream_batch_tokens_ = 1; VoxtralRealtimeTokenizer tokenizer_; VoxtralRealtimeFrontend frontend_; VoxtralRealtimeAudioEncoderRuntime audio_encoder_; @@ -62,6 +67,10 @@ class VoxtralRealtimeSession final runtime::AudioBuffer streaming_audio_; size_t streaming_audio_offset_values_ = 0; int64_t streaming_steps_processed_ = 0; + // Per-stage wall time summed over every steady step, reported by finalize(). + double stream_frontend_ms_ = 0.0; + double stream_encoder_ms_ = 0.0; + double stream_decoder_ms_ = 0.0; VoxtralRealtimeGenerationOptions streaming_generation_; runtime::StreamEventCallback stream_event_sink_; VoxtralRealtimeFrontendStreamState frontend_stream_state_; @@ -71,7 +80,6 @@ class VoxtralRealtimeSession final bool stream_started_ = false; bool first_stream_chunk_ = true; bool have_previous_stream_token_ = false; - bool stream_reached_eos_ = false; std::chrono::steady_clock::time_point stream_wall_start_{}; }; diff --git a/include/engine/models/voxtral_realtime/text_decoder.h b/include/engine/models/voxtral_realtime/text_decoder.h index 05682561..d3185393 100644 --- a/include/engine/models/voxtral_realtime/text_decoder.h +++ b/include/engine/models/voxtral_realtime/text_decoder.h @@ -20,7 +20,8 @@ class VoxtralRealtimeTextDecoderRuntime { size_t prefill_graph_arena_bytes, size_t decode_graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType weight_storage_type); + assets::TensorStorageType weight_storage_type, + int64_t stream_decode_cache_steps); ~VoxtralRealtimeTextDecoderRuntime(); VoxtralRealtimeGeneratedTokens generate( @@ -31,12 +32,17 @@ class VoxtralRealtimeTextDecoderRuntime { const VoxtralRealtimePrompt & prompt, const VoxtralRealtimeAudioEmbeddings & audio_embeddings, const VoxtralRealtimeGenerationOptions & options); + // `audio_row` selects one row of a possibly batched encoder pass; each call still advances + // the decoder by exactly one 80 ms step. int32_t stream_step( int32_t previous_token, const VoxtralRealtimeAudioEmbeddings & audio_embeddings, + int64_t audio_row, int64_t num_delay_tokens, const VoxtralRealtimeGenerationOptions & options); bool is_eos(int32_t token) const; + // Emits the summed upload/compute/readback/sample split of every streaming decode step. + void log_stream_timings() const; private: std::unique_ptr impl_; diff --git a/src/models/voxtral_realtime/audio_encoder.cpp b/src/models/voxtral_realtime/audio_encoder.cpp index 00954040..83a51c5c 100644 --- a/src/models/voxtral_realtime/audio_encoder.cpp +++ b/src/models/voxtral_realtime/audio_encoder.cpp @@ -93,10 +93,16 @@ assets::TensorStorageType normalize_weight_storage(assets::TensorStorageType sto case assets::TensorStorageType::F16: case assets::TensorStorageType::BF16: case assets::TensorStorageType::Q8_0: + // Requantized from the source GGUF at load time. Tensors whose rows are not a whole + // number of blocks (the 3-wide conv kernels) fall back to F32 in the weight store. + case assets::TensorStorageType::Q4_0: + case assets::TensorStorageType::Q4_K: + case assets::TensorStorageType::Q5_K: + case assets::TensorStorageType::Q6_K: return storage_type; default: throw std::runtime_error( - "VoxTral audio_encoder_weight_type supports only native, f32, f16, bf16, and q8_0"); + "VoxTral audio_encoder_weight_type supports only native, f32, f16, bf16, q4_0, q4_k, q5_k, q6_k, and q8_0"); } } @@ -219,8 +225,9 @@ core::TensorValue set_audio_kv_rows( const core::TensorValue & row_indices) { core::validate_rank_between(cache, 4, 4, "cache"); core::validate_rank_between(rows, 4, 4, "rows"); - if (cache.type != GGML_TYPE_F32 || rows.type != GGML_TYPE_F32) { - throw std::runtime_error("VoxTral audio streaming cache update requires f32 tensors"); + // ggml_set_rows converts on write, so an f16 cache still takes f32 rows straight from RoPE. + if ((cache.type != GGML_TYPE_F32 && cache.type != GGML_TYPE_F16) || rows.type != GGML_TYPE_F32) { + throw std::runtime_error("VoxTral audio streaming cache update requires f32 rows and an f32 or f16 cache"); } if (row_indices.type != GGML_TYPE_I32 || row_indices.shape.rank != 1 || row_indices.shape.dims[0] != rows.shape.dims[1]) { @@ -277,17 +284,18 @@ core::TensorValue build_audio_attention_streaming_static( q = rope.build(ctx, q, positions); k = rope.build(ctx, k, positions); q = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + q = core::ensure_backend_addressable_layout(ctx, q); k = core::ensure_backend_addressable_layout(ctx, k); v = core::ensure_backend_addressable_layout(ctx, v); auto updated_key = set_audio_kv_rows(ctx, key_cache, k, cache_slots); auto updated_value = set_audio_kv_rows(ctx, value_cache, v, cache_slots); + // Feed the sliding-window cache to flash attention as a strided view. Materializing the + // transpose would copy the full 750-step K and V caches once per layer per 80 ms chunk. auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, updated_key.shape.rank}).build(ctx, updated_key); auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, updated_value.shape.rank}).build(ctx, updated_value); - k_heads = core::wrap_tensor(ggml_cont(ctx.ggml, k_heads.tensor), k_heads.shape, k_heads.type); - v_heads = core::wrap_tensor(ggml_cont(ctx.ggml, v_heads.tensor), v_heads.shape, v_heads.type); auto context = modules::ScaledDotProductAttentionModule({ config.head_dim, - modules::ScaledDotProductAttentionLowering::Flash, + modules::ScaledDotProductAttentionLowering::FlashPreserveViews, GGML_PREC_F32, }).build(ctx, q, k_heads, v_heads, attention_mask); context = core::ensure_backend_addressable_layout(ctx, context); @@ -502,27 +510,30 @@ class StreamingAudioCache { conv2_cache_ = ggml_new_tensor_3d(ctx_.get(), GGML_TYPE_F32, 1, config.audio.hidden_size, 1); key_caches_.reserve(static_cast(config.audio.num_hidden_layers)); value_caches_.reserve(static_cast(config.audio.num_hidden_layers)); + // The sliding-window cache is read in full by flash attention on every 80 ms chunk and is + // never exported, so it is held in F16: half the read traffic, and the type Metal's flash + // attention kernels prefer. for (int64_t layer = 0; layer < config.audio.num_hidden_layers; ++layer) { key_caches_.push_back(core::wrap_tensor( ggml_new_tensor_4d( ctx_.get(), - GGML_TYPE_F32, + GGML_TYPE_F16, config.audio.head_dim, config.audio.num_key_value_heads, cache_steps_, 1), core::TensorShape::from_dims({1, cache_steps_, config.audio.num_key_value_heads, config.audio.head_dim}), - GGML_TYPE_F32)); + GGML_TYPE_F16)); value_caches_.push_back(core::wrap_tensor( ggml_new_tensor_4d( ctx_.get(), - GGML_TYPE_F32, + GGML_TYPE_F16, config.audio.head_dim, config.audio.num_key_value_heads, cache_steps_, 1), core::TensorShape::from_dims({1, cache_steps_, config.audio.num_key_value_heads, config.audio.head_dim}), - GGML_TYPE_F32)); + GGML_TYPE_F16)); } state_buffer_.reset(ggml_backend_alloc_ctx_tensors(ctx_.get(), backend_)); if (state_buffer_ == nullptr) { @@ -695,12 +706,17 @@ class StreamingAudioGraph { .build(ctx, x, {weights_->projector2_weight, std::nullopt}); output_ = x.tensor; ggml_set_output(output_); - ggml_set_output(next_conv1_cache_); - ggml_set_output(next_conv2_cache_); graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + // Expand the encoder body first so the conv-cache writes land after every node that reads + // the caches: the backend runs nodes in graph order, so this ordering is what keeps the + // in-graph update from clobbering the current chunk's own left context. ggml_build_forward_expand(graph_, output_); - ggml_build_forward_expand(graph_, next_conv1_cache_); - ggml_build_forward_expand(graph_, next_conv2_cache_); + ggml_build_forward_expand( + graph_, + ggml_cpy(ctx_.get(), next_conv1_cache_, cache_.conv1_cache())); + ggml_build_forward_expand( + graph_, + ggml_cpy(ctx_.get(), next_conv2_cache_, cache_.conv2_cache())); gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_))); if (gallocr_ == nullptr || !ggml_gallocr_reserve(gallocr_.get(), graph_) || @@ -728,6 +744,7 @@ class StreamingAudioGraph { if (features.frames != feature_frames_ || features.mel_bins != config.frontend.feature_size) { throw std::runtime_error("VoxTral streaming audio encoder feature shape mismatch"); } + const auto upload_start = Clock::now(); auto input_value = core::wrap_tensor(input_, core::TensorShape::from_dims({1, config.frontend.feature_size, feature_frames_}), GGML_TYPE_F32); core::write_tensor_f32(input_value, features.values); if (state.seen_encoder_steps == 0) { @@ -736,25 +753,36 @@ class StreamingAudioGraph { write_positions(state.seen_encoder_steps); write_cache_slots(state.seen_encoder_steps); write_mask(state.seen_encoder_steps); + upload_ms_ += engine::debug::elapsed_ms(upload_start); core::set_backend_threads(backend_, threads_); + const auto compute_start = Clock::now(); const ggml_status status = core::compute_backend_graph(backend_, graph_); ggml_backend_synchronize(backend_); + compute_ms_ += engine::debug::elapsed_ms(compute_start); if (status != GGML_STATUS_SUCCESS) { throw std::runtime_error("VoxTral streaming audio encoder graph compute failed"); } state.first_chunk = false; state.seen_encoder_steps += current_steps_; state.cached_encoder_steps = std::min(cache_steps_, state.cached_encoder_steps + current_steps_); - ggml_backend_tensor_copy_async(backend_, backend_, next_conv1_cache_, cache_.conv1_cache()); - ggml_backend_tensor_copy_async(backend_, backend_, next_conv2_cache_, cache_.conv2_cache()); - ggml_backend_synchronize(backend_); VoxtralRealtimeAudioEmbeddings out; + const auto readback_start = Clock::now(); out.values = core::read_tensor_f32(output_); + readback_ms_ += engine::debug::elapsed_ms(readback_start); out.tokens = tokens_; out.hidden_size = config.hidden_size; + ++runs_; return out; } + void log_timings(const char * prefix) const { + const std::string base = std::string("voxtral_realtime.audio_encoder.stream.") + prefix; + engine::debug::timing_log_scalar(base + ".runs", runs_); + engine::debug::timing_log_scalar(base + ".upload_ms", upload_ms_); + engine::debug::timing_log_scalar(base + ".compute_ms", compute_ms_); + engine::debug::timing_log_scalar(base + ".readback_ms", readback_ms_); + } + private: void write_positions(int64_t start) { std::vector positions(static_cast(current_steps_)); @@ -810,6 +838,10 @@ class StreamingAudioGraph { ggml_tensor * output_ = nullptr; std::vector mask_values_; ggml_cgraph * graph_ = nullptr; + int64_t runs_ = 0; + double upload_ms_ = 0.0; + double compute_ms_ = 0.0; + double readback_ms_ = 0.0; }; } // namespace @@ -884,6 +916,15 @@ struct VoxtralRealtimeAudioEncoderRuntime::Impl { return graph_slot->run(features, state); } + void log_stream_timings() const { + if (first_stream_graph != nullptr) { + first_stream_graph->log_timings("first"); + } + if (steady_stream_graph != nullptr) { + steady_stream_graph->log_timings("steady"); + } + } + std::shared_ptr assets; ggml_backend_t backend = nullptr; core::BackendType backend_type = core::BackendType::Cpu; @@ -925,4 +966,8 @@ VoxtralRealtimeAudioEmbeddings VoxtralRealtimeAudioEncoderRuntime::encode_stream return impl_->encode_stream_chunk(features, state); } +void VoxtralRealtimeAudioEncoderRuntime::log_stream_timings() const { + impl_->log_stream_timings(); +} + } // namespace engine::models::voxtral_realtime diff --git a/src/models/voxtral_realtime/frontend.cpp b/src/models/voxtral_realtime/frontend.cpp index c4ef28d3..fc91a088 100644 --- a/src/models/voxtral_realtime/frontend.cpp +++ b/src/models/voxtral_realtime/frontend.cpp @@ -250,7 +250,8 @@ VoxtralRealtimeFeatures VoxtralRealtimeFrontend::extract(const runtime::AudioBuf VoxtralRealtimeFeatures VoxtralRealtimeFrontend::extract_stream_chunk( const runtime::AudioBuffer & audio, bool first_chunk, - VoxtralRealtimeFrontendStreamState & state) const { + VoxtralRealtimeFrontendStreamState & state, + int64_t steady_tokens) const { const auto & config = assets_->config.frontend; if (audio.sample_rate <= 0 || audio.channels <= 0) { throw std::runtime_error("VoxTral audio input requires positive sample_rate and channels"); @@ -267,10 +268,10 @@ VoxtralRealtimeFeatures VoxtralRealtimeFrontend::extract_stream_chunk( mono.resize(static_cast(first_stream_chunk_samples()), 0.0F); mono = padded_first_stream_audio(assets_->config, std::move(mono)); } else { - if (steady_stream_chunk_samples() <= 0) { + if (steady_stream_chunk_samples(steady_tokens) <= 0) { throw std::runtime_error("VoxTral streaming frontend requires positive steady chunk samples"); } - mono.resize(static_cast(steady_stream_chunk_samples()), 0.0F); + mono.resize(static_cast(steady_stream_chunk_samples(steady_tokens)), 0.0F); } const bool reuse_cached_prefix = !first_chunk && state.cached_frame_ready; return extract_features_from_mono(config, mono, first_chunk, mel_filterbank_, &state, reuse_cached_prefix, false); @@ -282,9 +283,12 @@ int64_t VoxtralRealtimeFrontend::first_stream_chunk_samples() const { config.frontend.hop_length + config.frontend.win_length / 2; } -int64_t VoxtralRealtimeFrontend::steady_stream_chunk_samples() const { +int64_t VoxtralRealtimeFrontend::steady_stream_chunk_samples(int64_t steady_tokens) const { const auto & config = assets_->config; - return config.audio_length_per_tok * config.frontend.hop_length + config.frontend.win_length; + if (steady_tokens <= 0) { + throw std::runtime_error("VoxTral streaming frontend requires a positive steady token count"); + } + return steady_tokens * config.audio_length_per_tok * config.frontend.hop_length + config.frontend.win_length; } int64_t VoxtralRealtimeFrontend::first_stream_chunk_advance_samples() const { @@ -293,9 +297,12 @@ int64_t VoxtralRealtimeFrontend::first_stream_chunk_advance_samples() const { return frames * config.frontend.hop_length - config.frontend.n_fft / 2; } -int64_t VoxtralRealtimeFrontend::steady_stream_chunk_advance_samples() const { +int64_t VoxtralRealtimeFrontend::steady_stream_chunk_advance_samples(int64_t steady_tokens) const { const auto & config = assets_->config; - return config.audio_length_per_tok * config.frontend.hop_length; + if (steady_tokens <= 0) { + throw std::runtime_error("VoxTral streaming frontend requires a positive steady token count"); + } + return steady_tokens * config.audio_length_per_tok * config.frontend.hop_length; } } // namespace engine::models::voxtral_realtime diff --git a/src/models/voxtral_realtime/session.cpp b/src/models/voxtral_realtime/session.cpp index 7b79053b..6005842a 100644 --- a/src/models/voxtral_realtime/session.cpp +++ b/src/models/voxtral_realtime/session.cpp @@ -19,6 +19,12 @@ constexpr size_t kDefaultAudioEncoderWeightContextBytes = 128ull * 1024ull * 102 constexpr size_t kDefaultTextPrefillGraphArenaBytes = 512ull * 1024ull * 1024ull; constexpr size_t kDefaultTextDecodeGraphArenaBytes = 512ull * 1024ull * 1024ull; constexpr size_t kDefaultTextWeightContextBytes = 128ull * 1024ull * 1024ull; +// 1024 steps of 80 ms = ~82 s of decoder attention context, held in one graph for the whole +// session so a live stream never pays a cache-growth rebuild. +constexpr int64_t kDefaultStreamDecodeCacheSteps = 1024; +// Audio tokens per steady encoder forward. Batching amortizes the encoder's fixed per-dispatch +// cost across N tokens but delays every partial by N * 80 ms, so it stays opt-in. +constexpr int64_t kDefaultStreamBatchTokens = 1; int64_t source_frames_for_target_samples(int64_t target_samples, int source_rate, int channels, int target_rate) { if (target_samples <= 0 || source_rate <= 0 || channels <= 0 || target_rate <= 0) { @@ -102,16 +108,25 @@ VoxtralRealtimeSession::VoxtralRealtimeSession( } return weight_type; }()), + stream_decode_cache_steps_( + runtime::parse_i64_option(options.options, {"voxtral_realtime.stream_decode_cache_steps"}) + .value_or(kDefaultStreamDecodeCacheSteps)), + stream_batch_tokens_( + runtime::parse_i64_option(options.options, {"voxtral_realtime.stream_batch_tokens"}) + .value_or(kDefaultStreamBatchTokens)), tokenizer_(assets_), frontend_(assets_), audio_encoder_(assets_, execution_context(), audio_encoder_graph_arena_bytes_, audio_encoder_weight_context_bytes_, audio_encoder_weight_storage_type_), - text_decoder_(assets_, execution_context(), text_decoder_prefill_graph_arena_bytes_, text_decoder_decode_graph_arena_bytes_, text_decoder_weight_context_bytes_, text_decoder_weight_storage_type_) { + text_decoder_(assets_, execution_context(), text_decoder_prefill_graph_arena_bytes_, text_decoder_decode_graph_arena_bytes_, text_decoder_weight_context_bytes_, text_decoder_weight_storage_type_, stream_decode_cache_steps_) { if (task_.task != runtime::VoiceTaskKind::Asr) { throw std::runtime_error("VoxTral realtime only supports ASR"); } if (task_.mode != runtime::RunMode::Offline && task_.mode != runtime::RunMode::Streaming) { throw std::runtime_error("VoxTral realtime supports offline and streaming ASR"); } + if (stream_batch_tokens_ <= 0) { + throw std::runtime_error("VoxTral realtime stream_batch_tokens must be positive"); + } for (const auto & [key, value] : options.options) { (void) value; if (key.rfind("voxtral_realtime.", 0) == 0 && @@ -122,6 +137,8 @@ VoxtralRealtimeSession::VoxtralRealtimeSession( key != "voxtral_realtime.text_decoder_weight_context_mb" && key != "voxtral_realtime.audio_encoder_weight_type" && key != "voxtral_realtime.text_decoder_weight_type" && + key != "voxtral_realtime.stream_decode_cache_steps" && + key != "voxtral_realtime.stream_batch_tokens" && key != "voxtral_realtime.weight_type") { throw std::runtime_error("unknown VoxTral realtime session option: " + key); } @@ -224,8 +241,9 @@ runtime::StreamingPolicy VoxtralRealtimeSession::streaming_policy() const { runtime::StreamingPolicy policy; policy.input = runtime::StreamingInputKind::AudioChunks; policy.output = runtime::StreamingOutputKind::FinalResult; - policy.preferred_audio_chunk_samples = frontend_.steady_stream_chunk_samples(); - policy.preferred_audio_chunk_seconds = static_cast(frontend_.steady_stream_chunk_samples()) / + policy.preferred_audio_chunk_samples = frontend_.steady_stream_chunk_samples(stream_batch_tokens_); + policy.preferred_audio_chunk_seconds = + static_cast(frontend_.steady_stream_chunk_samples(stream_batch_tokens_)) / static_cast(assets_->config.frontend.sample_rate); return policy; } @@ -251,6 +269,9 @@ void VoxtralRealtimeSession::reset() { streaming_audio_ = runtime::AudioBuffer{}; streaming_audio_offset_values_ = 0; streaming_steps_processed_ = 0; + stream_frontend_ms_ = 0.0; + stream_encoder_ms_ = 0.0; + stream_decoder_ms_ = 0.0; streaming_generation_ = VoxtralRealtimeGenerationOptions{}; frontend_stream_state_ = VoxtralRealtimeFrontendStreamState{}; audio_stream_state_ = audio_encoder_.make_stream_state(); @@ -258,7 +279,6 @@ void VoxtralRealtimeSession::reset() { previous_stream_token_ = 0; first_stream_chunk_ = true; have_previous_stream_token_ = false; - stream_reached_eos_ = false; stream_started_ = false; stream_wall_start_ = {}; } @@ -318,6 +338,11 @@ runtime::TaskResult VoxtralRealtimeSession::finalize() { } engine::debug::timing_log_scalar("voxtral_realtime.session.stream.tokens", streaming_token_ids_.size()); engine::debug::timing_log_scalar("voxtral_realtime.session.stream.steps_processed", streaming_steps_processed_); + engine::debug::timing_log_scalar("voxtral_realtime.session.stream.frontend_ms", stream_frontend_ms_); + engine::debug::timing_log_scalar("voxtral_realtime.session.stream.encoder_ms", stream_encoder_ms_); + engine::debug::timing_log_scalar("voxtral_realtime.session.stream.decoder_ms", stream_decoder_ms_); + audio_encoder_.log_stream_timings(); + text_decoder_.log_stream_timings(); engine::debug::timing_log_scalar( "voxtral_realtime.session.stream.finalize_ms", engine::debug::elapsed_ms(finalize_start)); @@ -344,10 +369,12 @@ runtime::StreamEvent VoxtralRealtimeSession::process_available_stream_chunks() { throw std::runtime_error("VoxTral streaming pending audio has invalid channel layout"); } int64_t processed_chunks = 0; - while (!stream_reached_eos_) { + // Only the audio source running dry ends the loop: the decoder runs one step per 80 ms of + // audio for the whole session, so a pause in live audio never ends it. + while (true) { const int64_t target_samples = first_stream_chunk_ ? frontend_.first_stream_chunk_samples() - : frontend_.steady_stream_chunk_samples(); + : frontend_.steady_stream_chunk_samples(stream_batch_tokens_); const int64_t chunk_frames = source_frames_for_target_samples( target_samples, streaming_audio_.sample_rate, @@ -355,7 +382,7 @@ runtime::StreamEvent VoxtralRealtimeSession::process_available_stream_chunks() { static_cast(assets_->config.frontend.sample_rate)); const int64_t advance_target_samples = first_stream_chunk_ ? frontend_.first_stream_chunk_advance_samples() - : frontend_.steady_stream_chunk_advance_samples(); + : frontend_.steady_stream_chunk_advance_samples(stream_batch_tokens_); const int64_t advance_frames = source_frames_for_target_samples( advance_target_samples, streaming_audio_.sample_rate, @@ -379,11 +406,17 @@ runtime::StreamEvent VoxtralRealtimeSession::process_available_stream_chunks() { const auto begin = streaming_audio_.samples.begin() + static_cast(streaming_audio_offset_values_); chunk.samples.assign(begin, begin + static_cast(take_values)); streaming_audio_offset_values_ += advance_values; + const int64_t chunk_steps = first_stream_chunk_ ? 1 : stream_batch_tokens_; last_event = process_one_stream_chunk(chunk); - ++streaming_steps_processed_; + streaming_steps_processed_ += chunk_steps; ++processed_chunks; if (stream_event_sink_ != nullptr && last_event.partial_text.has_value()) { + // Several chunks can be consumed per process_audio_chunk call, so the sink is what + // delivers a partial per step rather than only the last one. The caller forwards + // whatever this function returns, so drop what the sink already delivered instead of + // emitting the same partial twice. stream_event_sink_(last_event); + last_event.partial_text.reset(); } if (pending_frames < chunk_frames) { break; @@ -407,35 +440,54 @@ runtime::StreamEvent VoxtralRealtimeSession::process_available_stream_chunks() { runtime::StreamEvent VoxtralRealtimeSession::process_one_stream_chunk(const runtime::AudioBuffer & audio) { runtime::StreamEvent event; event.is_final = false; - const auto features = frontend_.extract_stream_chunk(audio, first_stream_chunk_, frontend_stream_state_); + const auto frontend_start = Clock::now(); + const auto features = frontend_.extract_stream_chunk( + audio, + first_stream_chunk_, + frontend_stream_state_, + stream_batch_tokens_); + stream_frontend_ms_ += engine::debug::elapsed_ms(frontend_start); + const auto encoder_start = Clock::now(); const auto audio_embeddings = audio_encoder_.encode_stream_chunk(features, audio_stream_state_); - int32_t token = 0; + stream_encoder_ms_ += engine::debug::elapsed_ms(encoder_start); + const auto decoder_start = Clock::now(); if (first_stream_chunk_) { auto prompt = tokenizer_.build_transcription_prompt(frontend_.first_stream_chunk_samples(), true); - token = text_decoder_.begin_stream(prompt, audio_embeddings, streaming_generation_); + const int32_t token = text_decoder_.begin_stream(prompt, audio_embeddings, streaming_generation_); first_stream_chunk_ = false; - } else { - if (!have_previous_stream_token_) { - throw std::runtime_error("VoxTral streaming decoder is missing previous token"); - } - token = text_decoder_.stream_step( + stream_decoder_ms_ += engine::debug::elapsed_ms(decoder_start); + record_stream_token(token, event); + return event; + } + if (!have_previous_stream_token_) { + throw std::runtime_error("VoxTral streaming decoder is missing previous token"); + } + // One encoder forward covers `stream_batch_tokens_` audio tokens; the decoder still runs one + // step per token, each consuming its own row of the batch. + for (int64_t row = 0; row < audio_embeddings.tokens; ++row) { + const int32_t token = text_decoder_.stream_step( previous_stream_token_, audio_embeddings, + row, assets_->config.default_num_delay_tokens, streaming_generation_); + record_stream_token(token, event); } + stream_decoder_ms_ += engine::debug::elapsed_ms(decoder_start); + return event; +} + +void VoxtralRealtimeSession::record_stream_token(int32_t token, runtime::StreamEvent & event) { previous_stream_token_ = token; have_previous_stream_token_ = true; - if (text_decoder_.is_eos(token)) { - stream_reached_eos_ = true; - return event; - } + // The reference implementation gives EOS no special treatment: whatever token was sampled + // feeds back as the next text-stream input, and special ids simply emit no text. Only the + // audio source running dry ends the stream. if (!tokenizer_.is_stream_text_token(token)) { - return event; + return; } streaming_token_ids_.push_back(token); event.partial_text = runtime::Transcript{tokenizer_.decode(streaming_token_ids_), ""}; - return event; } } // namespace engine::models::voxtral_realtime diff --git a/src/models/voxtral_realtime/text_decoder.cpp b/src/models/voxtral_realtime/text_decoder.cpp index ddcc5781..6fdeea2a 100644 --- a/src/models/voxtral_realtime/text_decoder.cpp +++ b/src/models/voxtral_realtime/text_decoder.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -82,9 +83,16 @@ assets::TensorStorageType normalize_weight_storage(assets::TensorStorageType sto case assets::TensorStorageType::F16: case assets::TensorStorageType::BF16: case assets::TensorStorageType::Q8_0: + // Requantized from the source GGUF at load time. Tensors whose rows are not a whole + // number of blocks (the AdaRMS 32-wide projection) fall back to F32 in the weight store. + case assets::TensorStorageType::Q4_0: + case assets::TensorStorageType::Q4_K: + case assets::TensorStorageType::Q5_K: + case assets::TensorStorageType::Q6_K: return storage_type; default: - throw std::runtime_error("VoxTral text_decoder_weight_type supports only native, f32, f16, bf16, and q8_0"); + throw std::runtime_error( + "VoxTral text_decoder_weight_type supports only native, f32, f16, bf16, q4_0, q4_k, q5_k, q6_k, and q8_0"); } } @@ -305,6 +313,16 @@ struct TextAttentionOutput { core::TensorValue value; }; +// Summed across every streaming decode step, including the ones that ran on a graph that has +// since been replaced by a cache-growth rebuild. +struct DecodeStepStats { + int64_t steps = 0; + double upload_ms = 0.0; + double compute_ms = 0.0; + double readback_ms = 0.0; + double sample_ms = 0.0; +}; + TextAttentionOutput build_text_attention( core::ModuleBuildContext & ctx, const VoxtralRealtimeTextConfig & config, @@ -391,13 +409,13 @@ TextAttentionOutput build_text_attention_static( auto updated_value = set_rows.build(ctx, cache_value, v, cache_slot); auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); + // Feed the cache to flash attention as a strided view. Materializing the transpose would copy + // the whole K and V cache every layer, every token, which is the dominant cost at large caches. auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, updated_key.shape.rank}).build(ctx, updated_key); auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, updated_value.shape.rank}).build(ctx, updated_value); - k_heads = core::wrap_tensor(ggml_cont(ctx.ggml, k_heads.tensor), k_heads.shape, k_heads.type); - v_heads = core::wrap_tensor(ggml_cont(ctx.ggml, v_heads.tensor), v_heads.shape, v_heads.type); auto context = modules::GroupedQueryAttentionModule({ config.head_dim, - modules::GroupedQueryAttentionLowering::FlashGrouped, + modules::GroupedQueryAttentionLowering::FlashGroupedViewKV, GGML_PREC_F32, }).build(ctx, q_heads, k_heads, v_heads, attention_mask); context = core::ensure_backend_addressable_layout(ctx, context); @@ -823,17 +841,22 @@ class DecodeStepGraph { int32_t token, const VoxtralRealtimeAudioEmbeddings & audio_embeddings, int64_t num_delay_tokens, - bool single_audio_embedding, + // Which row of `audio_embeddings` this step consumes. Streaming passes the row explicitly + // because a batched encoder pass emits several rows; offline leaves it unset and indexes + // the embedding table by decode position. + std::optional audio_row, const VoxtralRealtimeGenerationOptions & options, uint64_t & sample_call_index, - TokenSampler & sampler) { + TokenSampler & sampler, + DecodeStepStats * stats = nullptr) { const auto & config = assets_->config.text; + const auto upload_start = Clock::now(); const int64_t position = step.position; if (position < 0) { throw std::runtime_error("VoxTral text decoder audio embedding position is out of range"); } - const int64_t audio_index = single_audio_embedding ? 0 : position; - if (audio_index >= audio_embeddings.tokens) { + const int64_t audio_index = audio_row.value_or(position); + if (audio_index < 0 || audio_index >= audio_embeddings.tokens) { throw std::runtime_error("VoxTral text decoder audio embedding position is out of range"); } ggml_backend_tensor_set(token_id_, &token, 0, sizeof(token)); @@ -856,14 +879,24 @@ class DecodeStepGraph { ggml_backend_tensor_set(cache_slot_, &step.cache_slot, 0, sizeof(step.cache_slot)); write_attention_mask(step); core::set_backend_threads(backend_, threads_); + const auto compute_start = Clock::now(); const ggml_status status = core::compute_backend_graph(backend_, graph_); ggml_backend_synchronize(backend_); + const auto readback_start = Clock::now(); if (status != GGML_STATUS_SUCCESS) { throw std::runtime_error("VoxTral text decoder step graph compute failed"); } logits_values_.resize(static_cast(config.vocab_size)); ggml_backend_tensor_get(logits_, logits_values_.data(), 0, logits_values_.size() * sizeof(float)); + const auto sample_start = Clock::now(); const int32_t next_token = sampler.select(logits_values_, options, options.seed, sample_call_index); + if (stats != nullptr) { + ++stats->steps; + stats->upload_ms += engine::debug::elapsed_ms(upload_start, compute_start); + stats->compute_ms += engine::debug::elapsed_ms(compute_start, readback_start); + stats->readback_ms += engine::debug::elapsed_ms(readback_start, sample_start); + stats->sample_ms += engine::debug::elapsed_ms(sample_start); + } return next_token; } @@ -930,23 +963,48 @@ struct VoxtralRealtimeTextDecoderRuntime::Impl { size_t prefill_graph_arena_bytes, size_t decode_graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + int64_t stream_decode_cache_steps) : assets(std::move(assets)), backend(execution.backend()), backend_type(execution.backend_type()), threads(std::max(1, execution.config().threads)), - graph_arena_bytes(std::max(prefill_graph_arena_bytes, decode_graph_arena_bytes)) { + graph_arena_bytes(std::max(prefill_graph_arena_bytes, decode_graph_arena_bytes)), + stream_decode_cache_steps(stream_decode_cache_steps) { if (this->assets == nullptr) { throw std::runtime_error("VoxTral text decoder requires assets"); } + if (this->stream_decode_cache_steps <= 0) { + throw std::runtime_error("VoxTral stream_decode_cache_steps must be positive"); + } + if (this->stream_decode_cache_steps > this->assets->config.text.sliding_window) { + throw std::runtime_error("VoxTral stream_decode_cache_steps exceeds the model sliding window"); + } weights = load_weights(*this->assets, backend, backend_type, weight_context_bytes, storage_type); + use_offline_decode_config(); + } + + // Growth buckets from min_cache_steps up to the sliding window: the offline path knows its + // token budget up front, so it can start small and grow at most a handful of times. + void use_offline_decode_config() { decode_runtime = runtime::BoundedStaticKVDecodeRuntime({ - this->assets->config.text.sliding_window, + assets->config.text.sliding_window, 256, "VoxTral text decoder bounded KV", }); } + // A live stream has no token budget, so pinning both bounds to the same value makes + // target_cache_steps constant: the graph is built once at begin_stream and never rebuilt, + // trading the full 8192-step window for the absence of mid-stream rebuild stalls. + void use_stream_decode_config() { + decode_runtime = runtime::BoundedStaticKVDecodeRuntime({ + stream_decode_cache_steps, + stream_decode_cache_steps, + "VoxTral text decoder streaming KV", + }); + } + VoxtralRealtimeGeneratedTokens generate( const VoxtralRealtimePrompt & prompt, const VoxtralRealtimeAudioEmbeddings & audio_embeddings, @@ -962,6 +1020,7 @@ struct VoxtralRealtimeTextDecoderRuntime::Impl { return {}; } const auto total_start = Clock::now(); + use_offline_decode_config(); uint64_t sample_call_index = 0; const int64_t prompt_steps = static_cast(prompt.input_ids.size()); const auto decode_graph_factory = [this](int64_t cache_steps) { @@ -1024,7 +1083,7 @@ struct VoxtralRealtimeTextDecoderRuntime::Impl { token, audio_embeddings, prompt.num_delay_tokens, - false, + std::nullopt, options, sample_call_index, sampler); @@ -1056,7 +1115,9 @@ struct VoxtralRealtimeTextDecoderRuntime::Impl { throw std::runtime_error("VoxTral streaming text decoder prompt/audio token count mismatch"); } const auto total_start = Clock::now(); + use_stream_decode_config(); stream_sample_call_index_ = 0; + stream_stats_ = DecodeStepStats{}; const auto decode_graph_factory = [this](int64_t cache_steps) { auto graph = std::make_unique( assets, @@ -1111,13 +1172,14 @@ struct VoxtralRealtimeTextDecoderRuntime::Impl { int32_t stream_step( int32_t previous_token, const VoxtralRealtimeAudioEmbeddings & audio_embeddings, + int64_t audio_row, int64_t num_delay_tokens, const VoxtralRealtimeGenerationOptions & options) { if (!decode_runtime.has_graph()) { throw std::runtime_error("VoxTral streaming text decoder requires begin_stream before stream_step"); } - if (audio_embeddings.tokens != 1) { - throw std::runtime_error("VoxTral streaming text decoder expects one audio token per steady chunk"); + if (audio_row < 0 || audio_row >= audio_embeddings.tokens) { + throw std::runtime_error("VoxTral streaming text decoder audio row is out of range"); } const auto decode_graph_factory = [this](int64_t cache_steps) { auto graph = std::make_unique( @@ -1138,10 +1200,11 @@ struct VoxtralRealtimeTextDecoderRuntime::Impl { previous_token, audio_embeddings, num_delay_tokens, - true, + audio_row, options, stream_sample_call_index_, - sampler); + sampler, + &stream_stats_); decode_runtime.advance_after_direct_append(1); return token; } @@ -1150,16 +1213,26 @@ struct VoxtralRealtimeTextDecoderRuntime::Impl { return token == assets->config.text.eos_token_id; } + void log_stream_timings() const { + engine::debug::timing_log_scalar("voxtral_realtime.text_decoder.stream.steps", stream_stats_.steps); + engine::debug::timing_log_scalar("voxtral_realtime.text_decoder.stream.upload_ms", stream_stats_.upload_ms); + engine::debug::timing_log_scalar("voxtral_realtime.text_decoder.stream.compute_ms", stream_stats_.compute_ms); + engine::debug::timing_log_scalar("voxtral_realtime.text_decoder.stream.readback_ms", stream_stats_.readback_ms); + engine::debug::timing_log_scalar("voxtral_realtime.text_decoder.stream.sample_ms", stream_stats_.sample_ms); + } + std::shared_ptr assets; ggml_backend_t backend = nullptr; core::BackendType backend_type = core::BackendType::Cpu; int threads = 1; size_t graph_arena_bytes = 0; + int64_t stream_decode_cache_steps = 0; std::shared_ptr weights; std::unique_ptr prefill_graph; runtime::BoundedStaticKVDecodeRuntime decode_runtime; TokenSampler sampler; uint64_t stream_sample_call_index_ = 0; + DecodeStepStats stream_stats_; }; VoxtralRealtimeTextDecoderRuntime::VoxtralRealtimeTextDecoderRuntime( @@ -1168,14 +1241,16 @@ VoxtralRealtimeTextDecoderRuntime::VoxtralRealtimeTextDecoderRuntime( size_t prefill_graph_arena_bytes, size_t decode_graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType weight_storage_type) + assets::TensorStorageType weight_storage_type, + int64_t stream_decode_cache_steps) : impl_(std::make_unique( std::move(assets), execution, prefill_graph_arena_bytes, decode_graph_arena_bytes, weight_context_bytes, - weight_storage_type)) {} + weight_storage_type, + stream_decode_cache_steps)) {} VoxtralRealtimeTextDecoderRuntime::~VoxtralRealtimeTextDecoderRuntime() = default; @@ -1196,13 +1271,18 @@ int32_t VoxtralRealtimeTextDecoderRuntime::begin_stream( int32_t VoxtralRealtimeTextDecoderRuntime::stream_step( int32_t previous_token, const VoxtralRealtimeAudioEmbeddings & audio_embeddings, + int64_t audio_row, int64_t num_delay_tokens, const VoxtralRealtimeGenerationOptions & options) { - return impl_->stream_step(previous_token, audio_embeddings, num_delay_tokens, options); + return impl_->stream_step(previous_token, audio_embeddings, audio_row, num_delay_tokens, options); } bool VoxtralRealtimeTextDecoderRuntime::is_eos(int32_t token) const { return impl_->is_eos(token); } +void VoxtralRealtimeTextDecoderRuntime::log_stream_timings() const { + impl_->log_stream_timings(); +} + } // namespace engine::models::voxtral_realtime diff --git a/tests/unittests/test_voxtral_realtime_stream_chunking.cpp b/tests/unittests/test_voxtral_realtime_stream_chunking.cpp new file mode 100644 index 00000000..0e623b08 --- /dev/null +++ b/tests/unittests/test_voxtral_realtime_stream_chunking.cpp @@ -0,0 +1,144 @@ +#include "engine/models/voxtral_realtime/frontend.h" + +#include "test_assert.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using engine::models::voxtral_realtime::VoxtralRealtimeAssets; +using engine::models::voxtral_realtime::VoxtralRealtimeFrontend; +using engine::models::voxtral_realtime::VoxtralRealtimeFrontendStreamState; +using engine::test::require; +using engine::test::require_eq; + +// The frontend only needs the config to size chunks and build its mel filterbank, so the whole +// suite runs without a model, weights or a backend. +VoxtralRealtimeFrontend make_frontend() { + auto assets = std::make_shared(); + return VoxtralRealtimeFrontend(std::move(assets)); +} + +engine::runtime::AudioBuffer mono_audio(int64_t samples) { + engine::runtime::AudioBuffer audio; + audio.sample_rate = 16000; + audio.channels = 1; + audio.samples.assign(static_cast(samples), 0.25F); + return audio; +} + +bool throws(const std::function & call) { + try { + call(); + } catch (const std::runtime_error &) { + return true; + } + return false; +} + +// A steady chunk holds `steady_tokens` audio tokens of 8 * 160 = 1280 samples (80 ms), plus one +// win_length tail shared by the whole batch. Batching must not change what a single token costs. +void test_steady_chunk_samples_scale_per_token() { + const auto frontend = make_frontend(); + require_eq(frontend.steady_stream_chunk_samples(1), 1680, "steady chunk samples at 1 token"); + for (const int64_t tokens : {1, 2, 4, 7}) { + require_eq( + frontend.steady_stream_chunk_samples(tokens), + tokens * 1280 + 400, + "steady chunk samples at " + std::to_string(tokens) + " tokens"); + } +} + +// The advance is what the session consumes per chunk, and it is what +// `streaming_steps_processed_ += stream_batch_tokens_` is accounted against: a batched chunk must +// advance by exactly as much as the same number of single-token chunks would. +void test_steady_advance_is_linear_in_tokens() { + const auto frontend = make_frontend(); + const int64_t single = frontend.steady_stream_chunk_advance_samples(1); + require_eq(single, 1280, "steady advance at 1 token"); + for (const int64_t tokens : {1, 2, 4, 7}) { + require_eq( + frontend.steady_stream_chunk_advance_samples(tokens), + tokens * single, + "steady advance at " + std::to_string(tokens) + " tokens"); + } +} + +// Every caller that does not batch must be bit-for-bit unaffected by the new parameter. +void test_defaults_match_single_token() { + const auto frontend = make_frontend(); + require_eq(frontend.steady_stream_chunk_samples(), 1680, "default steady chunk samples"); + require_eq(frontend.steady_stream_chunk_advance_samples(), 1280, "default steady advance"); + // The first chunk is never batched: it primes the decoder with the delay-token lookahead. + require_eq(frontend.first_stream_chunk_samples(), 9000, "first chunk samples"); + require_eq(frontend.first_stream_chunk_advance_samples(), 8760, "first chunk advance"); +} + +void test_non_positive_token_counts_are_rejected() { + const auto frontend = make_frontend(); + for (const int64_t tokens : {static_cast(0), static_cast(-1)}) { + require( + throws([&] { (void) frontend.steady_stream_chunk_samples(tokens); }), + "steady_stream_chunk_samples must reject " + std::to_string(tokens)); + require( + throws([&] { (void) frontend.steady_stream_chunk_advance_samples(tokens); }), + "steady_stream_chunk_advance_samples must reject " + std::to_string(tokens)); + require( + throws([&] { + VoxtralRealtimeFrontendStreamState state; + (void) frontend.extract_stream_chunk(mono_audio(1680), false, state, tokens); + }), + "extract_stream_chunk must reject " + std::to_string(tokens)); + } +} + +// One STFT pass over an N-token chunk must yield 8N feature frames, which the encoder's conv +// stack turns into 4N encoder steps and, after downsample_factor 4, exactly N audio tokens. The +// first steady chunk and the ones after it take different STFT routes -- the first computes the +// whole magnitude, later ones reuse the previous chunk's cached frame as their prefix -- so both +// are exercised here and must agree. +void test_batched_chunk_yields_one_token_worth_of_frames_each() { + const auto frontend = make_frontend(); + for (const int64_t tokens : {1, 2, 4}) { + VoxtralRealtimeFrontendStreamState state; + const auto audio = mono_audio(frontend.steady_stream_chunk_samples(tokens)); + const std::string label = " at " + std::to_string(tokens) + " tokens"; + + const auto uncached = frontend.extract_stream_chunk(audio, false, state, tokens); + require_eq(uncached.frames, tokens * 8, "uncached steady frames" + label); + require_eq(uncached.mel_bins, 128, "uncached steady mel bins" + label); + require(state.cached_frame_ready, "steady chunk must cache its STFT tail" + label); + + const auto cached = frontend.extract_stream_chunk(audio, false, state, tokens); + require_eq(cached.frames, uncached.frames, "cached steady frames" + label); + require_eq(cached.mel_bins, uncached.mel_bins, "cached steady mel bins" + label); + require_eq( + static_cast(cached.values.size()), + cached.frames * cached.mel_bins, + "cached steady feature size" + label); + } +} + +} // namespace + +int main() { + try { + test_steady_chunk_samples_scale_per_token(); + test_steady_advance_is_linear_in_tokens(); + test_defaults_match_single_token(); + test_non_positive_token_counts_are_rejected(); + test_batched_chunk_yields_one_token_worth_of_frames_each(); + std::cout << "voxtral_realtime_stream_chunking_test passed\n"; + } catch (const std::exception & ex) { + std::cerr << "voxtral_realtime_stream_chunking_test failed: " << ex.what() << "\n"; + return 1; + } + return 0; +}