diff --git a/conversion/asr.py b/conversion/asr.py index 08702fd7..9e4d2512 100644 --- a/conversion/asr.py +++ b/conversion/asr.py @@ -69,7 +69,8 @@ re.compile(r"^joint\.enc\.weight$"), re.compile(r"^joint\.pred\.weight$"), re.compile(r"^joint\.joint_net\.2\.weight$"), - # ConformerConv pointwise convs are 1×1 — mathematically Linear. We squeeze + # ConformerConv pointwise convs are 1×1 — mathematically Linear. Cache-aware + # RNNT Q8 exports retain F16 here; see _pick_dtype. We squeeze # the k=1 dim in reshape_for_emission so the Q8_0/k-quant kernels see a # 2D (out, in) tensor; the runtime detects 2D stored shape and dispatches # the matmul through ggml_mul_mat (see Conv1D::build_graph). @@ -132,11 +133,18 @@ def _q8_0_to_tensor_planar(packed: np.ndarray, n_per_row: int) -> np.ndarray: def _pick_dtype( - name: str, shape: tuple, linear_qtype: GGMLQuantizationType, default_qtype: GGMLQuantizationType + name: str, + shape: tuple, + linear_qtype: GGMLQuantizationType, + default_qtype: GGMLQuantizationType, + *, + cache_aware_rnnt: bool = False, ) -> tuple[GGMLQuantizationType, Optional[str]]: """Return (chosen GGMLQuantizationType, fallback reason or None). Dispatch by tensor class: + * Cache-aware RNNT Q8 pointwise convs → F16; quantizing these kernels + can cause long non-speech/blank runs in speech. * Linear weights → user-selected `--weight-type` (linear_qtype). * Conv weights, embed → F16 (regardless of --weight-type) — saves ~50% disk vs F32 with no runtime cost. @@ -144,6 +152,12 @@ def _pick_dtype( pos_bias, pe table, etc. The runtime keeps these in F32 in memory. """ + if ( + cache_aware_rnnt + and linear_qtype == GGMLQuantizationType.Q8_0 + and _POINTWISE_PATTERN.fullmatch(name) + ): + return GGMLQuantizationType.F16, "cache-aware RNNT pointwise precision" if _is_linear_weight(name): inner = int(shape[-1]) if len(shape) >= 1 else 1 if linear_qtype in _K_QUANTS and inner % _QK_K != 0: @@ -677,7 +691,11 @@ def convert( arr = reshape_for_emission(dst_name, arr) chosen_qtype, fallback_reason = _pick_dtype( - dst_name, arr.shape, linear_qtype, default_qtype + dst_name, + arr.shape, + linear_qtype, + default_qtype, + cache_aware_rnnt=cache_supported and head_type == "rnnt", ) if fallback_reason is not None: print(f"[convert] kept fp16 ({fallback_reason}): {dst_name}") diff --git a/docs/model-conversion.md b/docs/model-conversion.md index 0cc16112..4964cd86 100644 --- a/docs/model-conversion.md +++ b/docs/model-conversion.md @@ -61,3 +61,13 @@ option list. The converters do not import `nemo_toolkit`, and source checkpoints are not modified. + +### Cache-aware RNNT Q8 precision + +Cache-aware RNNT Q8_0 exports keep the Conformer pointwise convolution weights +in F16. Quantizing these kernels to Q8 can cause speech to be decoded as long +blank runs; attention, feed-forward and decoder matrix quantization is unchanged. +Existing GGUFs are not repaired by updating the runtime: regenerate them from +the original `.nemo` checkpoint with `--outtype q8_0`. Converting already-quantized +weights to F16 cannot recover the original values. Other model/head types and +requested weight types keep their previous conversion policy. diff --git a/src/asr/encoder/cache_aware_encoder.h b/src/asr/encoder/cache_aware_encoder.h index 0278c4d4..d78ead7a 100644 --- a/src/asr/encoder/cache_aware_encoder.h +++ b/src/asr/encoder/cache_aware_encoder.h @@ -1,6 +1,6 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -// Cache-aware streaming encoder: one shared, fixed-shape FastConformer encoder +// Cache-aware streaming encoder: one shared FastConformer encoder // Session whose per-layer K/V caches are two planes of one persistent arena // and whose convolution cache is a persistent matrix, with one indexed row per // stream. A microbatch supplies stream IDs, gathers those rows, and scatters @@ -65,8 +65,9 @@ class CacheAwareEncoder { int chunk_frames(); // One cache-aware encoder step for the stream owning `state`. `mel` is - // (n_mels x n_mel_frames) frame-major (= pre_encode_cache_size + sub*(1+R) - // frames). The stream's K/V/conv cache lives in `state` (device-resident, + // (n_mels x n_mel_frames) frame-major: first chunk = 1+sub*R; subsequent + // chunks = pre_encode_cache_size+sub*(1+R). Both produce 1+R encoder frames. + // The stream's K/V/conv cache lives in `state` (device-resident, // updated in-graph). `enc_out` receives (output_dim x T_enc) flat-packed // (d_model without a tail); T_enc = 1 + R. Thread-safe. void encode( diff --git a/src/asr/encoder/fastconformer.cpp b/src/asr/encoder/fastconformer.cpp index 08e9497d..29063974 100644 --- a/src/asr/encoder/fastconformer.cpp +++ b/src/asr/encoder/fastconformer.cpp @@ -1410,11 +1410,17 @@ FastConformerEncoder::build_graph( pre.set_first_tensor(ggml_runtime::ggml_bf_tensor(x_scaled, x.buft)); } + // Only subsequent chunks carry overlap. The shorter first-chunk shape + // already produces exactly cache_chunk_frames, with causal padding handled + // independently by each subsampling convolution. A shape-based distinction + // also keeps concurrent first/steady-state streams in separate graph keys. + const bool first_chunk = + input_tensors.get_tensor(0).tensor->ne[1] == cfg_.cache_first_chunk_mel_frames(); // Then drop the leading `cache_drop_extra` frames from the subsampled // output so the visible chunk length is exactly cache_chunk_frames // (= 1 + R). The dropped frames overlap with the tail of the previous // chunk (handled by the K/V cache). - if (cfg_.cache_drop_extra > 0) { + if (!first_chunk && cfg_.cache_drop_extra > 0) { auto x = pre.get_tensor(0); auto bf_ctx = session_tensor_container->get_ctx_of_buffer_type(x.buft); auto x_view = ggml_view_4d( diff --git a/src/asr/encoder/fastconformer.h b/src/asr/encoder/fastconformer.h index e21b8605..bad28f9a 100644 --- a/src/asr/encoder/fastconformer.h +++ b/src/asr/encoder/fastconformer.h @@ -14,6 +14,7 @@ #pragma once #include +#include #include #include @@ -72,6 +73,16 @@ struct EncoderConfig { // 2 frames, giving cache_chunk_frames = 1 + R after the drop. int cache_drop_extra = 2; + // Causal subsampling starts at frame zero, without a preceding overlap. + // Subsequent-chunk geometry remains owned by the runner. + int cache_first_chunk_mel_frames() const { return 1 + subsampling_factor * cache_right_ctx; } + + // Exclusive mel frontier: causal encoder frame f ends at mel f*sub. + // The first emitted frame therefore covers one real mel frame, not sub. + int64_t cache_decoded_mel_frames(int64_t emitted_frames) const { + return emitted_frames > 0 ? 1 + (emitted_frames - 1) * subsampling_factor : 0; + } + // Offline attention limits are independent of the streaming cache context. // -1 means unlimited; finite values preserve a model's trained attention window. int offline_left_ctx = -1; diff --git a/src/asr/runner.cpp b/src/asr/runner.cpp index d13b25d9..74c63ca6 100644 --- a/src/asr/runner.cpp +++ b/src/asr/runner.cpp @@ -874,12 +874,20 @@ CacheStreamRunner::feed_audio(const float* samples, size_t n_samples) { audio_buf_.insert(audio_buf_.end(), samples, samples + n_samples); } +int +CacheStreamRunner::next_chunk_mel_frames() const { + // Preserve the existing overlap and steady-state formula; only startup + // uses the shorter causal chunk. + return first_chunk_ ? enc_cfg_.cache_first_chunk_mel_frames() + : pre_encode_cache_size_ + + enc_cfg_.subsampling_factor * (1 + enc_cfg_.cache_right_ctx); +} + void CacheStreamRunner::process_one_chunk() { - // NeMo formula: chunk_size_mel = pre_encode_cache_size + sub * (1 + R). const int sub = enc_cfg_.subsampling_factor; const int R = enc_cfg_.cache_right_ctx; - const int chunk_size_mel = pre_encode_cache_size_ + sub * (1 + R); + const int chunk_size_mel = next_chunk_mel_frames(); const int shift_size_mel = sub * (1 + R - cache_drop_size_); const int n_mels = model_->fe_config().n_mels; const int left_ctx = enc_cfg_.cache_left_ctx; @@ -945,8 +953,23 @@ CacheStreamRunner::process_one_chunk() { // Advance mel buffer: drop shift_size_mel; keep overlap_mel_frames at // the head so the next chunk's leading frames see the same context. - const size_t shift_floats = static_cast(shift_size_mel) * n_mels; - mel_offset_ = std::min(mel_offset_ + shift_floats, mel_buf_.size()); + if (first_chunk_) { + first_chunk_ = false; + if (chunk_size_mel < pre_encode_cache_size_) { + // R=0 begins with one real mel frame. Its first subsequent chunk + // still needs the existing overlap; only the unavailable negative + // history is zero, and only AFTER the first chunk was encoded. + compact_mel_buffer(); + mel_buf_.insert( + mel_buf_.begin(), + static_cast(pre_encode_cache_size_ - chunk_size_mel) * n_mels, 0.0f); + } else { + mel_offset_ += static_cast(chunk_size_mel - pre_encode_cache_size_) * n_mels; + } + } else { + const size_t shift_floats = static_cast(shift_size_mel) * n_mels; + mel_offset_ = std::min(mel_offset_ + shift_floats, mel_buf_.size()); + } } StreamingUpdate @@ -993,24 +1016,12 @@ CacheStreamRunner::step() { total_mel_frames_produced_ = i_start + n; } - const int sub = enc_cfg_.subsampling_factor; - const int R = enc_cfg_.cache_right_ctx; - const int chunk_size_mel = pre_encode_cache_size_ + sub * (1 + R); - // Process every whole chunk available, polling the endpointer per chunk // so a silence gap interior to one large feed still fires. An EOU breaks // the loop; the next utterance's remaining chunks decode on the next // step() call. last_step_new_tokens_.clear(); - while (static_cast((mel_buf_.size() - mel_offset_) / n_mels) >= chunk_size_mel) { - if (!stream_zero_padded_) { - // First-chunk zero-pad so the encoder's unconditional - // cache_drop_extra eats zero-pad, not the leading 160 ms of real audio. - compact_mel_buffer(); - mel_buf_.insert( - mel_buf_.begin(), static_cast(pre_encode_cache_size_) * n_mels, 0.0f); - stream_zero_padded_ = true; - } + while (static_cast((mel_buf_.size() - mel_offset_) / n_mels) >= next_chunk_mel_frames()) { process_one_chunk(); if (poll_endpoint(update, /*after_chunk=*/true)) break; @@ -1051,9 +1062,10 @@ CacheStreamRunner::poll_endpoint(StreamingUpdate& update, bool after_chunk) { double now_ms = 0.0, last_speech_ms = 0.0; if (vad_ && endpointer_->config().vad_based) { // VAD-driven: scan the Silero speech bits over the newly decoded mel - // frames (enc frame f covers mel frames [f*sub, (f+1)*sub)). + // frames. Causal encoder frame f ends at mel f*sub (inclusive), so + // the shorter first chunk shifts the exclusive frontier by sub-1. const double mel_ms = 1000.0 * hop / sample_rate; - const int64_t decoded_mel = total_frames_emitted_ * enc_cfg_.subsampling_factor; + const int64_t decoded_mel = enc_cfg_.cache_decoded_mel_frames(total_frames_emitted_); scan_vad_speech(*vad_, decoded_mel, vad_scan_frame_, vad_speech_seen_frame_); now_ms = static_cast(decoded_mel) * mel_ms; last_speech_ms = (vad_speech_seen_frame_ < 0) @@ -1105,7 +1117,7 @@ CacheStreamRunner::trim_buffers() { // masker's neighbourhood scans never cross the trimmed base. const int64_t margin = std::max(256, vad_masker_ ? 2 * vad_masker_->max_scan() : 0); - const int64_t decoded_mel = total_frames_emitted_ * enc_cfg_.subsampling_factor; + const int64_t decoded_mel = enc_cfg_.cache_decoded_mel_frames(total_frames_emitted_); const int64_t g = decoded_mel - margin; if (g > 0) vad_->discard_timeline_before(g); @@ -1153,15 +1165,6 @@ CacheStreamRunner::finalize() { const int shift_size_mel = sub * (1 + R - cache_drop_size_); const size_t buffered = (mel_buf_.size() - mel_offset_) / n_mels; if (buffered > 0) { - // A stream shorter than one chunk never entered step()'s chunk - // loop, so the start-of-stream zero-pad hasn't happened yet; apply - // it here or the encoder's cache_drop_extra eats the leading - // ~160 ms of real audio instead of pad. - if (!stream_zero_padded_) { - mel_buf_.insert( - mel_buf_.begin(), static_cast(pre_encode_cache_size_) * n_mels, 0.0f); - stream_zero_padded_ = true; - } // Tail flush: pad enough zero mel for the final real frame to // see a full right-context window AND for the RNNT predictor // to get at least one all-zero enc frame to commit pending tail tokens. @@ -1169,7 +1172,8 @@ CacheStreamRunner::finalize() { mel_buf_.resize(mel_buf_.size() + flush_pad_frames * static_cast(n_mels), 0.0f); const size_t before = all_tokens_.size(); last_step_new_tokens_.clear(); - while (static_cast((mel_buf_.size() - mel_offset_) / n_mels) >= chunk_size_mel) { + while (static_cast((mel_buf_.size() - mel_offset_) / n_mels) >= + next_chunk_mel_frames()) { process_one_chunk(); } for (size_t i = before; i < all_tokens_.size(); i++) { @@ -1217,7 +1221,7 @@ CacheStreamRunner::reset() { last_enc_T_ = 0; finalized_ = false; finalizing_ = false; - stream_zero_padded_ = false; + first_chunk_ = true; zero_caches(); std::fill(attn_mask_.begin(), attn_mask_.end(), 0.0f); if (head_) diff --git a/src/asr/runner.h b/src/asr/runner.h index cd45ffd3..dc33e224 100644 --- a/src/asr/runner.h +++ b/src/asr/runner.h @@ -243,8 +243,8 @@ class OfflineRunner final : public AsrRunner { bool exceeds_offline_position_limit(const AsrModel& model, size_t n_samples, int input_sample_rate); -// Drives a fixed-shape encoder graph (chunk_size_mel in, chunk_enc_frames -// out) with per-layer K/V/conv caches threaded across chunks in a +// Drives first/steady-state encoder shapes (each producing chunk_enc_frames) +// with per-layer K/V/conv caches threaded across chunks in a // device-resident indexed cache row. Audio is buffered until at least one chunk's // worth of mel frames is ready; each step uploads mel + a small validity mask // and downloads only the joint-projected encoder output. @@ -279,6 +279,7 @@ class CacheStreamRunner final : public AsrRunner { RnntDecodeStats rnnt_decode_stats() const; private: + int next_chunk_mel_frames() const; void process_one_chunk(); void upload_attn_mask(); void zero_caches(); @@ -297,7 +298,8 @@ class CacheStreamRunner final : public AsrRunner { // GGUF-pinned cache-aware mel-frame geometry. These are NeMo's // CacheAwareStreamingConfig chunk_size / shift_size in mel frames // (setup_streaming_params, conformer_encoder.py): - // chunk_size_mel = pre_encode_cache_size + subsampling_factor*(1 + R) + // first_chunk_mel = 1 + subsampling_factor*R + // subsequent_chunk_mel = pre_encode_cache_size + subsampling_factor*(1 + R) // shift_size_mel = subsampling_factor*(1 + R - cache_drop_size) // These values are fixed by the model. int pre_encode_cache_size_ = 9; // overlap, mel frames @@ -359,11 +361,9 @@ class CacheStreamRunner final : public AsrRunner { bool finalized_ = false; // True once finalize() begins, so its internal drain step() doesn't fire EOU. bool finalizing_ = false; - // Tracks whether we've prepended pre_encode_cache_size zero mel frames at - // start-of-stream. NeMo's CacheAwareStreamingAudioBuffer does this once - // per stream so the encoder's drop_extra_pre_encoded behaviour throws - // away zero-pad rather than the leading ~160 ms of real audio. - bool stream_zero_padded_ = false; + // Reset at an encoder-state boundary, independently of the absolute stream + // frame clock. The first chunk has no pre-encoder overlap to discard. + bool first_chunk_ = true; // Encoder output from most recent chunk (column-major (d_model, T_out)). std::vector last_enc_out_; diff --git a/tests/conversion/converter_contract_test.py b/tests/conversion/converter_contract_test.py index 60a21ad7..0fbeffe1 100644 --- a/tests/conversion/converter_contract_test.py +++ b/tests/conversion/converter_contract_test.py @@ -11,18 +11,22 @@ import tarfile import tempfile import unittest +from contextlib import redirect_stdout from pathlib import Path from unittest import mock import numpy as np import torch import yaml -from gguf import GGMLQuantizationType +from gguf import GGMLQuantizationType, GGUFReader from safetensors.torch import save_file ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(ROOT)) +from conversion.asr import _pick_dtype +from conversion.asr import convert as convert_asr +from conversion.asr import reshape_for_emission from conversion.diarization import is_high_resolution_v3 from conversion.diarization import remap as remap_diarization from conversion.registry import ( @@ -52,6 +56,93 @@ class ConverterContractTest(unittest.TestCase): + def test_streaming_rnnt_q8_preserves_pointwise_precision(self) -> None: + for layer in (0, 23): + for conv in (1, 2): + name = f"encoder.layers.{layer}.conv.pointwise_conv{conv}.weight" + shape = reshape_for_emission(name, np.zeros((64, 32, 1))).shape + self.assertEqual(shape, (64, 32)) + dtype, reason = _pick_dtype( + name, + shape, + GGMLQuantizationType.Q8_0, + GGMLQuantizationType.F32, + cache_aware_rnnt=True, + ) + self.assertEqual(dtype, GGMLQuantizationType.F16) + self.assertEqual(reason, "cache-aware RNNT pointwise precision") + # Offline/CTC and other export precisions retain their policy. + dtype, _ = _pick_dtype( + name, shape, GGMLQuantizationType.Q8_0, GGMLQuantizationType.F32 + ) + self.assertEqual(dtype, GGMLQuantizationType.Q8_0) + for requested in (GGMLQuantizationType.F16, GGMLQuantizationType.BF16): + dtype, _ = _pick_dtype( + name, + shape, + requested, + GGMLQuantizationType.F32, + cache_aware_rnnt=True, + ) + self.assertEqual(dtype, requested) + + def test_streaming_rnnt_precision_change_is_limited_to_pointwise(self) -> None: + expected = { + "encoder.layers.0.self_attn.linear_q.weight": GGMLQuantizationType.Q8_0, + "encoder.layers.0.feed_forward1.linear1.weight": GGMLQuantizationType.Q8_0, + "decoder.prediction.dec_rnn.lstm.ih_l0.weight": GGMLQuantizationType.Q8_0, + "encoder.layers.0.conv.depthwise_conv.weight": GGMLQuantizationType.F16, + "encoder.pre_encode.conv.0.weight": GGMLQuantizationType.F16, + "encoder.layers.0.conv.pointwise_conv1.bias": GGMLQuantizationType.F32, + } + for name, want in expected.items(): + dtype, reason = _pick_dtype( + name, + (64, 32), + GGMLQuantizationType.Q8_0, + GGMLQuantizationType.F32, + cache_aware_rnnt=True, + ) + self.assertEqual(dtype, want, name) + self.assertIsNone(reason) + + def test_streaming_rnnt_export_keeps_original_pointwise_values(self) -> None: + name = "encoder.layers.0.conv.pointwise_conv1.weight" + linear = "encoder.layers.0.feed_forward1.linear1.weight" + values = torch.arange(32 * 32, dtype=torch.float32).reshape(32, 32, 1) / 337 + config = { + "encoder": { + "d_model": 32, + "n_layers": 1, + "n_heads": 4, + "conv_kernel_size": 3, + "feat_in": 128, + "pos_emb_max_len": 32, + "att_context_size": [[70, 3]], + "att_context_style": "chunked_limited", + }, + "decoder": {"vocab_size": 32}, + } + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + torch.save({name: values, linear: values.squeeze(-1)}, root / "model_weights.ckpt") + for streaming in (True, False): + config["encoder"]["att_context_style"] = ( + "chunked_limited" if streaming else "regular" + ) + (root / "model_config.yaml").write_text(yaml.safe_dump(config), encoding="utf-8") + output = root / f"model-{streaming}.gguf" + with redirect_stdout(io.StringIO()): + convert_asr(root, output, "rnnt", "q8_0") + tensors = {tensor.name: tensor for tensor in GGUFReader(output).tensors} + self.assertEqual(tensors[linear].tensor_type, GGMLQuantizationType.Q8_0) + want = GGMLQuantizationType.F16 if streaming else GGMLQuantizationType.Q8_0 + self.assertEqual(tensors[name].tensor_type, want) + if streaming: + np.testing.assert_array_equal( + tensors[name].data, values.squeeze(-1).half().numpy() + ) + def _checkpoint(self, root: Path, name: str, config: dict) -> Path: checkpoint = root / f"{name}.nemo" payload = yaml.safe_dump(config).encode("utf-8") diff --git a/tests/cpp/asr/CMakeLists.txt b/tests/cpp/asr/CMakeLists.txt index 1eed7350..4f072841 100644 --- a/tests/cpp/asr/CMakeLists.txt +++ b/tests/cpp/asr/CMakeLists.txt @@ -9,6 +9,12 @@ target_link_libraries(test_ctc_streaming PRIVATE nemo_speech_asr) add_executable(test_rnnt_streaming test_rnnt_streaming.cpp) target_link_libraries(test_rnnt_streaming PRIVATE nemo_speech_asr) +add_executable(test_cache_stream_geometry test_cache_stream_geometry.cpp) +target_link_libraries(test_cache_stream_geometry PRIVATE nemo_speech_asr) +# Keep this separate from TEST_ASR_MODEL, which can also point to a CTC model. +add_test(NAME cache_stream_geometry + COMMAND test_cache_stream_geometry $ENV{NEMO_SPEECH_TEST_CACHE_AWARE_RNNT_MODEL}) + add_executable(test_gguf_mmap_loader test_gguf_mmap_loader.cpp) target_link_libraries(test_gguf_mmap_loader PRIVATE nemo_speech_asr) # Synthetic-GGUF cases always run; NEMO_SPEECH_TEST_ASR_MODEL (if set) adds a full-model compare. diff --git a/tests/cpp/asr/test_cache_stream_geometry.cpp b/tests/cpp/asr/test_cache_stream_geometry.cpp new file mode 100644 index 00000000..db967d26 --- /dev/null +++ b/tests/cpp/asr/test_cache_stream_geometry.cpp @@ -0,0 +1,252 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// No arguments: model-independent causal first/steady-state shape checks. +// Optional: MODEL.gguf [--gpu N] exercises actual streaming, reset and EOF. +#include +#include +#include +#include +#include +#include + +#include "recognizer.h" + +using namespace nemo_speech::asr; + +static void +require(bool condition, const char* message) { + if (!condition) + throw std::runtime_error(message); +} + +static void +check_shapes() { + for (int right : {0, 1, 3, 6, 13}) { + const auto cfg = make_cache_aware_config(EncoderConfig{}, right); + const int first = 1 + 8 * right; + const int steady = 9 + 8 * (1 + right); + require(cfg.cache_first_chunk_mel_frames() == first, "wrong first mel length"); + require(cfg.subsample_time_length(first) == 1 + right, "first chunk needs no drop"); + require( + cfg.subsample_time_length(steady) - cfg.cache_drop_extra == 1 + right, + "steady chunk must drop only the overlap"); + require(cfg.cache_decoded_mel_frames(0) == 0, "empty stream has a nonzero decode clock"); + require(cfg.cache_decoded_mel_frames(1) == 1, "first encoder frame must end at mel zero"); + for (int chunks : {1, 2, 100}) { + const int64_t emitted = static_cast(chunks) * (1 + right); + const int64_t consumed = first + static_cast(chunks - 1) * 8 * (1 + right); + require(cfg.cache_decoded_mel_frames(emitted) == consumed, "VAD decode clock is ahead"); + } + } + + // R13: 14 emitted frames cover 105 real mel frames, not 112. At a 10ms + // hop, the old clock would fire an 800ms endpoint after only 730ms silence. + const auto cfg = make_cache_aware_config(EncoderConfig{}, 13); + VadEndpointer endpoint(VadEndpointerCfg{true, true, 800.0f}); + require(!endpoint.poll(320.0, 320.0), "speech must not immediately endpoint"); + require( + !endpoint.poll(cfg.cache_decoded_mel_frames(14) * 10.0, 320.0), + "VAD endpoint fired on the old 70ms-ahead clock"); + require( + endpoint.poll(cfg.cache_decoded_mel_frames(15) * 10.0, 320.0), + "VAD endpoint did not fire after decoded silence reached the threshold"); +} + +static void +same_encoder(const std::vector& a, const std::vector& b) { + require(a.size() == b.size(), "encoder size differs from explicit reference chunk"); + float largest = 0.0f; + double sum = 0.0; + for (size_t i = 0; i < a.size(); ++i) { + require(std::isfinite(a[i]) && std::isfinite(b[i]), "nonfinite encoder projection"); + largest = std::max(largest, std::fabs(a[i] - b[i])); + sum += static_cast(a[i] - b[i]) * (a[i] - b[i]); + } + if (largest > 1e-5f) + std::fprintf( + stderr, "encoder max abs error=%g RMS=%g\n", largest, std::sqrt(sum / a.size())); + require(largest <= 1e-5f, "encoder differs from explicit first/history chunk"); +} + +static void +check_model(const std::string& path, int gpu, int right) { + RecognizerConfig cfg; + cfg.model.path = path; + cfg.backend.gpu = gpu; + cfg.streaming.rnnt_right_context = right; + cfg.log_status = false; + Recognizer recognizer(cfg); + auto* model = dynamic_cast(recognizer.model()); + require(model != nullptr, "expected cache-aware RNNT model"); + const auto enc = make_cache_aware_config(model->encoder_config(), right); + const auto& fe = model->fe_config(); + const int first_mel = 1 + enc.subsampling_factor * right; + const int new_mel = enc.subsampling_factor * (1 + right); + const int overlap = 9; // Existing runner policy, not inferred from subsampling. + const int hop = model->fe().hop_length(); + const size_t first_samples = (first_mel - 1) * hop + fe.n_fft / 2; + const size_t stride_samples = new_mel * hop; + std::vector audio(first_samples + 2 * stride_samples); + for (size_t i = 0; i < audio.size(); ++i) + audio[i] = 0.08f * std::sin(0.071f * i) + 0.03f * std::cos(0.019f * i); + + const int prompt = model->prompt_index_for_lang("auto"); + CacheStreamRunner runner(model, cfg); + runner.set_prompt_index(prompt); + runner.feed_audio(audio.data(), first_samples - 1); + runner.step(); + require(runner.chunks_processed() == 0, "first chunk ran before its last stable mel frame"); + runner.feed_audio(audio.data() + first_samples - 1, 1); + runner.step(); + require( + runner.chunks_processed() == 1, "first chunk incorrectly waits for steady-state length"); + + // Independent chunk assembly, including negative history for R=0. Use + // another cache slot so runner buffer movement cannot affect the oracle. + std::vector all_mel; + std::vector initial_audio(audio.begin(), audio.begin() + first_samples - 1); + produce_new_mel_frames(model->fe(), initial_audio, 0, 0, all_mel); + auto state = model->make_cache_state(); + int consumed = 0; + int cached = 0; + std::vector first_output; + for (int step = 0; step < 3; ++step) { + std::fprintf(stderr, "[check] R=%d chunk=%d\n", right, step); + // Match the available audio frontier, not future audio: the first + // reflected STFT frame of a minimal R=0 packet must use the same + // boundary in the runner and in this chunk-assembly reference. + const size_t available = first_samples + step * stride_samples; + std::vector prefix_audio(audio.begin(), audio.begin() + available); + std::vector new_features; + produce_new_mel_frames( + model->fe(), prefix_audio, 0, static_cast(all_mel.size() / fe.n_mels), + new_features); + all_mel.insert(all_mel.end(), new_features.begin(), new_features.end()); + const int prefix = step == 0 ? 0 : overlap; + const int fresh = step == 0 ? first_mel : new_mel; + const int frames = prefix + fresh; + std::vector chunk(static_cast(frames) * fe.n_mels, 0.0f); + for (int f = 0; f < frames; ++f) { + const int source = consumed - prefix + f; + if (source >= 0) { + require( + static_cast(source + 1) * fe.n_mels <= all_mel.size(), + "reference mel underrun"); + std::copy_n( + all_mel.data() + static_cast(source) * fe.n_mels, fe.n_mels, + chunk.data() + static_cast(f) * fe.n_mels); + } + } + std::vector mask(enc.cache_left_ctx + enc.cache_chunk_frames, 0.0f); + std::fill_n(mask.begin(), std::max(0, enc.cache_left_ctx - cached), -1e9f); + std::vector expected; + int expected_t = 0; + model->encode_cache_aware( + state, chunk.data(), frames, mask.data(), static_cast(mask.size()), expected, + expected_t, prompt); + require(expected_t == right + 1, "wrong reference encoder frame count"); + if (step > 0) { + runner.feed_audio( + audio.data() + first_samples + (step - 1) * stride_samples, stride_samples); + runner.step(); + } + std::vector actual; + int actual_t = 0, dim = 0; + runner.take_encoder_frames(actual, actual_t, dim); + require(actual_t == expected_t, "wrong runner encoder frame count"); + require(runner.chunks_processed() == step + 1, "wrong chunk consumption"); + same_encoder(actual, expected); + if (step == 0) + first_output = actual; + consumed += fresh; + cached = std::min(enc.cache_left_ctx, cached + expected_t); + } + runner.reset(); + // Preserve FE packet shapes as well: Metal FFT rounding can differ across + // packet shapes independently of the encoder's chunk geometry. + runner.feed_audio(audio.data(), first_samples - 1); + runner.step(); + runner.feed_audio(audio.data() + first_samples - 1, 1); + runner.step(); + require(runner.chunks_processed() == 1, "reset did not restore first-chunk geometry"); + std::vector reset_output; + int reset_t = 0, reset_dim = 0; + runner.take_encoder_frames(reset_output, reset_t, reset_dim); + same_encoder(first_output, reset_output); + + auto decode = [&](size_t size, size_t packet, bool drain) { + runner.reset(); + std::vector ids; + for (size_t off = 0; off < size;) { + const size_t count = std::min(packet, size - off); + runner.feed_audio(audio.data() + off, count); + off += count; + if (drain) { + const auto update = runner.step(); + ids.insert(ids.end(), update.new_token_ids.begin(), update.new_token_ids.end()); + } + } + const auto final = runner.finalize(); + ids.insert(ids.end(), final.new_token_ids.begin(), final.new_token_ids.end()); + require(final.is_final, "EOF not final"); + require(runner.finalize().new_token_ids.empty(), "EOF duplicated tokens"); + return ids; + }; + for (size_t size : + {size_t(0), size_t(1), size_t(159), first_samples - 1, first_samples, first_samples + 1, + audio.size()}) { + const auto expected = decode(size, audio.size(), false); + require(decode(size, 1279, true) == expected, "odd packet EOF differs"); + require(decode(size, 1600, true) == expected, "100ms packet EOF differs"); + } + // EOU publishes text without ending the acoustic stream. Preserve the + // first chunk's cache and require steady-state geometry after the event. + // Use identical FE packet shapes for the EOU and no-EOU controls: otherwise + // Metal FFT rounding differences would test packetization, not the boundary. + auto after_boundary = [&](bool force) { + runner.reset(); + runner.feed_audio(audio.data(), first_samples - 1); + runner.step(); + runner.feed_audio(audio.data() + first_samples - 1, 1); + if (force) + runner.force_eou(); + require(runner.step().is_final == force, "unexpected forced EOU result"); + require(runner.cache_filled_frames() == right + 1, "EOU changed encoder cache"); + const int before = runner.chunks_processed(); + require(before == 1, "EOU synthesized extra chunks"); + runner.feed_audio(audio.data() + first_samples, stride_samples - 1); + runner.step(); + require(runner.chunks_processed() == before, "post-EOU reused first-chunk geometry"); + runner.feed_audio(audio.data() + first_samples + stride_samples - 1, 1); + runner.step(); + require(runner.chunks_processed() == before + 1, "post-EOU steady chunk was not processed"); + std::vector output; + int count = 0, dim = 0; + runner.take_encoder_frames(output, count, dim); + return output; + }; + const auto no_eou = after_boundary(false); + same_encoder(no_eou, after_boundary(true)); + std::fprintf(stdout, "[PASS] R=%d first/history encoder, reset, EOF and EOU\n", right); +} + +int +main(int argc, char** argv) try { + check_shapes(); + if (argc > 1) { + int gpu = -1; + if (argc == 4 && std::string(argv[2]) == "--gpu") + gpu = std::stoi(argv[3]); + else if (argc != 2) + throw std::runtime_error("Usage: test_cache_stream_geometry [MODEL.gguf [--gpu N]]"); + for (int right : {0, 3, 6, 13}) check_model(argv[1], gpu, right); + } else { + std::fprintf(stdout, "[SKIP] model checks: no cache-aware RNNT model supplied\n"); + } + std::fprintf(stdout, "[PASS] cache-aware chunk geometry\n"); + return 0; +} +catch (const std::exception& error) { + std::fprintf(stderr, "[FAIL] %s\n", error.what()); + return 1; +} diff --git a/tests/cpp/asr/test_endpointer.cpp b/tests/cpp/asr/test_endpointer.cpp index 02858f3d..9e4a4acd 100644 --- a/tests/cpp/asr/test_endpointer.cpp +++ b/tests/cpp/asr/test_endpointer.cpp @@ -10,6 +10,8 @@ // utterance] clip, stream it through the runner with endpointing enabled, // and assert >= 2 finals with the right split. Token-silence mode by // default; pass --vad-model to also exercise the VAD-driven path. +// Cache-aware RNNT also checks each VAD endpoint against an independent +// first/steady mel-coverage timeline, before and after reset. // // Usage: ./test_endpointer [ [--gpu N] [--vad-model F] // [--chunk-ms N] [--gap-ms N] [--eou-ms N] @@ -113,6 +115,67 @@ test_policy() { // --------------------------------------------------------------------------- // 2. Integration: stream a 2-utterance clip, count finals. // --------------------------------------------------------------------------- +static void +test_cache_vad_clock( + RnntModel* model, const RecognizerConfig& cfg, const std::vector& audio) { + CacheStreamRunner stream(model, cfg); + stream.set_prompt_index(model->prompt_index_for_lang("auto")); + SileroVad reference(model->backend_manager(), cfg.vad.model_path); + const auto enc = + make_cache_aware_config(model->encoder_config(), cfg.streaming.rnnt_right_context); + const int hop = model->fe().hop_length(); + const int sub = enc.subsampling_factor; + const int first = 1 + sub * enc.cache_right_ctx; + const int stride = sub * (1 + enc.cache_right_ctx); + const double mel_ms = 1000.0 * hop / model->fe_config().sample_rate; + // At most one encoder chunk becomes ready per feed, including R0. The + // independent VAD sees the same packets, including not-yet-decoded audio. + const size_t packet = static_cast(sub) * hop; + reference.set_binarizer(cfg.vad.masker.onset, cfg.vad.masker.offset, hop); + for (int pass = 0; pass < 2; ++pass) { + stream.reset(); + reference.reset(); + VadEndpointer expected(cfg.endpointing); + int64_t scan = 0, speech = -1; + int endpoints = 0; + bool matched = true; + for (size_t off = 0; off < audio.size(); off += packet) { + const size_t count = std::min(packet, audio.size() - off); + reference.observe_audio(audio.data() + off, count); + const int before = stream.chunks_processed(); + stream.feed_audio(audio.data() + off, count); + const auto update = stream.step(); + const int chunks = stream.chunks_processed(); + if (chunks - before > 1) { + matched = false; + break; + } + // Derive coverage from real first/steady chunk lengths, not the + // production encoder-frame-to-mel helper being tested. + const int64_t decoded = chunks ? first + static_cast(chunks - 1) * stride : 0; + const int64_t end = std::min(decoded, reference.committed_frames()); + for (; scan < end; ++scan) + if (reference.frame_speech(scan)) + speech = scan; + const bool final = expected.poll(decoded * mel_ms, (speech + 1) * mel_ms); + endpoints += final ? 1 : 0; + if (update.is_final != final) { + std::fprintf( + stderr, + "VAD clock mismatch: pass=%d R=%d chunks=%d decoded_mel=%lld speech=%lld " + "expected_final=%d actual_final=%d\n", + pass, enc.cache_right_ctx, chunks, static_cast(decoded), + static_cast(speech), final, update.is_final); + matched = false; + break; + } + } + check(matched, "integration: VAD finals match real decoded-mel coverage, including reset"); + check( + speech >= 0 && endpoints > 0, "integration: VAD clock control sees speech and an EOU"); + } +} + static int test_integration(int argc, char** argv) { const std::string model_path = argv[1]; @@ -148,6 +211,10 @@ test_integration(int argc, char** argv) { p.pe_bin_path = const_cast(""); ggml_runtime::BackendManager bm(p); auto model = AsrModel::load(bm, model_path); + // Match production language detection for prompt-conditioned RNNT models. + const int prompt = model->head_kind() == HeadKind::Rnnt + ? static_cast(model.get())->prompt_index_for_lang("auto") + : -1; std::vector one; int sr = 0; @@ -175,12 +242,16 @@ test_integration(int argc, char** argv) { cfg.endpointing = ep; cfg.streaming.rnnt_right_context = right_ctx; // the CTC runner ignores it + if (model->head_kind() == HeadKind::Rnnt && !vad_model.empty()) + test_cache_vad_clock(static_cast(model.get()), cfg, audio); + std::unique_ptr runner; if (model->head_kind() == HeadKind::Rnnt) { runner = std::make_unique(static_cast(model.get()), cfg); } else { runner = std::make_unique(static_cast(model.get()), cfg); } + runner->set_prompt_index(prompt); AsrRequestOptions request_options; request_options.enable_word_time_offsets = true; request_options.enable_automatic_punctuation = punctuation; @@ -196,6 +267,7 @@ test_integration(int argc, char** argv) { if (model->head_kind() == HeadKind::Rnnt) { CacheStreamRunner silent(static_cast(model.get()), cfg); + silent.set_prompt_index(prompt); silent.set_request_options(request_options); std::vector silence(static_cast(3) * sr, 0.0f); const auto began = std::chrono::steady_clock::now(); @@ -231,6 +303,7 @@ test_integration(int argc, char** argv) { probe = std::make_unique(static_cast(model.get()), pcfg); } + probe->set_prompt_index(prompt); const size_t head = std::min(audio.size(), static_cast(4) * sr); const auto start = std::chrono::steady_clock::now(); for (size_t off = 0; off < head; off += chunk_samples) { @@ -307,6 +380,7 @@ test_integration(int argc, char** argv) { RecognizerConfig control = cfg; control.endpointing.enable = endpointing; CacheStreamRunner stream(static_cast(model.get()), control); + stream.set_prompt_index(prompt); stream.set_request_options(request_options); std::vector tokens; int endpoints = 0;