Bucket the Parakeet audio encoder

cactus-engine/src/engine.h

@@ -669,6 +669,8 @@ class ToolCallConstrainer {
     void feed_gemma_char(char ch);
 };
 
+size_t select_bucket(const std::vector<size_t>& capacities, size_t required_frames);
+
 class Model {
 public:
     Model();
@@ -1008,6 +1010,7 @@ class Model {
     int output_index(const Component& comp, const std::string& name) const;
     uint32_t argmax_last_logits(float* out_uncertainty = nullptr, float repetition_penalty = 1.0f);
     bool load_handoff_probe();
+    Component* select_audio_encoder(size_t required_frames);
     void maybe_capture_handoff_probe_hidden(const Component& comp, const std::string& output_name = "probe_hidden");
     void run_vision_encoder(const std::string& image_path);
     void run_vision_encoder_lfm2_vl(const std::string& image_path);

cactus-engine/src/model.cpp

@@ -4619,6 +4619,31 @@ std::vector<uint32_t> Model::transcribe_whisper_seq2seq(
         should_stop);
 }
 
+size_t select_bucket(const std::vector<size_t>& capacities, size_t required_frames) {
+    size_t best = capacities.size();
+    for (size_t i = 0; i < capacities.size(); ++i) {
+        const size_t capacity = capacities[i];
+        if (capacity == 0 || capacity < required_frames) continue;
+        if (best == capacities.size() || capacity < capacities[best]) best = i;
+    }
+    return best;
+}
+
+Model::Component* Model::select_audio_encoder(size_t required_frames) {
+    std::vector<Component*> candidates;
+    std::vector<size_t> capacities;
+    for (auto& [name, comp] : components_) {
+        if (name.rfind("audio_encoder", 0) != 0) continue;
+        auto it = comp.metadata.find("audio_frames");
+        if (it == comp.metadata.end()) continue;
+        candidates.push_back(&comp);
+        capacities.push_back(parse_size_or_zero(it->second));
+    }
+    const size_t index = select_bucket(capacities, required_frames);
+    if (index < candidates.size()) return candidates[index];
+    return components_.count("audio_encoder") ? &components_.at("audio_encoder") : nullptr;
+}
+
 std::vector<uint32_t> Model::transcribe_parakeet_tdt(const std::vector<float>& audio_features,
                                                      ParakeetTdtStreamState* stream, bool is_final,
                                                      size_t end_frame,
@@ -4628,7 +4653,11 @@ std::vector<uint32_t> Model::transcribe_parakeet_tdt(const std::vector<float>& a
 
     reset_handoff_probe_rollout();
 
-    Component* audio_enc = components_.count("audio_encoder") ? &components_.at("audio_encoder") : nullptr;
+    const size_t config_mels = static_cast<size_t>(config_.num_mel_bins);
+    const size_t required_frames = config_mels > 0
+        ? audio_features.size() / config_mels
+        : std::numeric_limits<size_t>::max();
+    Component* audio_enc = select_audio_encoder(required_frames);
     Component* dec = components_.count("decoder") ? &components_.at("decoder") : nullptr;
     if (!audio_enc || !dec) {
         CACTUS_LOG_ERROR("model", "Parakeet TDT bundle missing audio_encoder or decoder component");

cactus-engine/tests/test_bucket.cpp

@@ -0,0 +1,78 @@
+#include "test_utils.h"
+#include "../src/engine.h"
+
+#include <limits>
+#include <vector>
+
+using namespace TestUtils;
+using cactus::engine::select_bucket;
+
+namespace {
+
+const std::vector<size_t> kFibonacci = {100, 200, 300, 500, 800, 1300, 2003};
+
+bool picks_exact_fit() {
+    return select_bucket(kFibonacci, 300) == 2;
+}
+
+bool picks_smallest_that_fits() {
+    return select_bucket(kFibonacci, 201) == 2
+        && select_bucket(kFibonacci, 301) == 3
+        && select_bucket(kFibonacci, 801) == 5;
+}
+
+bool picks_first_bucket_for_tiny_input() {
+    return select_bucket(kFibonacci, 1) == 0
+        && select_bucket(kFibonacci, 100) == 0;
+}
+
+bool reports_none_when_above_every_bucket() {
+    return select_bucket(kFibonacci, 2004) == kFibonacci.size();
+}
+
+bool reports_none_for_empty_ladder() {
+    return select_bucket({}, 100) == 0;
+}
+
+bool ignores_zero_capacities() {
+    const std::vector<size_t> with_zero = {0, 500, 0, 200};
+    return select_bucket(with_zero, 150) == 3
+        && select_bucket(with_zero, 300) == 1;
+}
+
+bool is_independent_of_declaration_order() {
+    const std::vector<size_t> shuffled = {2003, 300, 100, 1300, 200, 800, 500};
+    return shuffled[select_bucket(shuffled, 250)] == 300
+        && shuffled[select_bucket(shuffled, 900)] == 1300;
+}
+
+bool handles_zero_required_frames() {
+    return select_bucket(kFibonacci, 0) == 0;
+}
+
+bool unknown_required_frames_falls_back() {
+    return select_bucket(kFibonacci, std::numeric_limits<size_t>::max()) == kFibonacci.size();
+}
+
+bool handles_duplicate_capacities() {
+    const std::vector<size_t> duplicated = {300, 300, 800};
+    return duplicated[select_bucket(duplicated, 250)] == 300;
+}
+
+}
+
+int main() {
+    TestRunner runner("Audio Encoder Bucket Selection");
+    runner.run_test("Exact fit", picks_exact_fit());
+    runner.run_test("Smallest that fits", picks_smallest_that_fits());
+    runner.run_test("Tiny input takes first bucket", picks_first_bucket_for_tiny_input());
+    runner.run_test("None when above every bucket", reports_none_when_above_every_bucket());
+    runner.run_test("None for empty ladder", reports_none_for_empty_ladder());
+    runner.run_test("Zero capacities ignored", ignores_zero_capacities());
+    runner.run_test("Order independent", is_independent_of_declaration_order());
+    runner.run_test("Zero required frames", handles_zero_required_frames());
+    runner.run_test("Unknown required frames falls back", unknown_required_frames_falls_back());
+    runner.run_test("Duplicate capacities", handles_duplicate_capacities());
+    runner.print_summary();
+    return runner.all_passed() ? 0 : 1;
+}

cactus-engine/tests/test_llm.cpp

@@ -740,23 +740,17 @@ bool test_decode_batch_throughput() {
 
 int main() {
     TestUtils::TestRunner runner("LLM Tests");
-    runner.run_test("1k_context", test_1k_context());
     runner.run_test("streaming", test_streaming());
     runner.run_test("prefill", test_prefill());
-    runner.run_test("prefill_idempotent_reuse", test_prefill_idempotent_reuse());
-    runner.run_test("prefill_prefix_extension_reuse", test_prefill_prefix_extension_reuse());
     runner.run_test("prefill_invalidated_on_message_change", test_prefill_invalidated_on_message_change());
     runner.run_test("chunked_prefill_padding", test_chunked_prefill_padding());
     runner.run_test("tool_calls", test_tool_call());
     runner.run_test("tool_multiple_tool_call_invocations", test_multiple_tool_call_invocations());
-    runner.run_test("tool_calls_with_three_tools", test_tool_call_with_three_tools());
     runner.run_test("tool_constraint_clear_releases_bias", test_tool_constraint_clear_releases_bias());
     runner.run_test("partition_thinking_response", test_partition_thinking_response());
     runner.run_test("prompt_retains_thinking", test_prompt_gemma4_retains_thinking());
     runner.run_test("complete_thinking_api_clean", test_complete_gemma4_thinking_api_clean());
     runner.run_test("multiturn_thinking_persist", test_multiturn_thinking_persist());
-    runner.run_test("multiturn_turn2_distinct", test_multiturn_turn2_distinct());
-    runner.run_test("decode_batch", test_decode_batch());
     runner.run_test("generate_batch_ragged", test_generate_batch_ragged());
     runner.run_test("decode_batch_throughput", test_decode_batch_throughput());
     runner.run_test("batch_distinct4_matches_single", test_batch_distinct4_matches_single());

python/cactus/convert/tests/test_audio_buckets.py

@@ -0,0 +1,50 @@
+from cactus.transpile.audio_preprocess import audio_bucket_frames
+from cactus.transpile.audio_preprocess import audio_bucket_seconds
+
+
+def test_fibonacci_ladder_capped_at_model_window():
+    assert audio_bucket_frames(2003) == [100, 200, 300, 500, 800, 1300, 2003]
+
+
+def test_window_below_ladder_yields_single_bucket():
+    assert audio_bucket_frames(80) == [80]
+
+
+def test_window_equal_to_a_rung_is_not_duplicated():
+    assert audio_bucket_frames(800) == [100, 200, 300, 500, 800]
+
+
+def test_non_positive_window_yields_no_buckets():
+    assert audio_bucket_frames(0) == []
+    assert audio_bucket_frames(-1) == []
+
+
+def test_buckets_are_sorted_and_unique():
+    buckets = audio_bucket_frames(2003)
+    assert buckets == sorted(buckets)
+    assert len(buckets) == len(set(buckets))
+
+
+def test_default_ladder_is_fibonacci_seconds():
+    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13)
+
+
+def test_env_override_replaces_ladder(monkeypatch):
+    monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "2,4,6")
+    assert audio_bucket_seconds() == (2, 4, 6)
+    assert audio_bucket_frames(2003) == [200, 400, 600, 2003]
+
+
+def test_env_override_is_sorted_and_deduplicated(monkeypatch):
+    monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "6, 2, 4, 2")
+    assert audio_bucket_seconds() == (2, 4, 6)
+
+
+def test_malformed_env_override_falls_back_to_default(monkeypatch):
+    monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "not-a-number")
+    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13)
+
+
+def test_empty_env_override_falls_back_to_default(monkeypatch):
+    monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "0,-3")
+    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13)

python/cactus/transpile/audio_preprocess.py

@@ -21,6 +21,7 @@
 _PARAKEET_PREEMPHASIS = 0.97
 _PARAKEET_LOG_FLOOR = np.float32(2**-24)
 _DEFAULT_MAX_AUDIO_SECONDS = 30.0
+_DEFAULT_AUDIO_BUCKET_SECONDS = (1, 2, 3, 5, 8, 13)
 
 
 def audio_duration_limit_seconds() -> float:
@@ -32,6 +33,33 @@ def audio_duration_limit_seconds() -> float:
     return max(0.0, value)
 
 
+def audio_bucket_seconds() -> tuple[int, ...]:
+    raw = os.environ.get("CACTUS_TRANSPILER_AUDIO_BUCKETS")
+    if not raw:
+        return _DEFAULT_AUDIO_BUCKET_SECONDS
+    try:
+        values = sorted({int(part) for part in raw.split(",")})
+    except (TypeError, ValueError):
+        return _DEFAULT_AUDIO_BUCKET_SECONDS
+    positive = tuple(value for value in values if value > 0)
+    if not positive:
+        return _DEFAULT_AUDIO_BUCKET_SECONDS
+    return positive
+
+
+def audio_bucket_frames(max_frames: int) -> list[int]:
+    if max_frames <= 0:
+        return []
+    frames_per_second = _PARAKEET_SAMPLE_RATE // _PARAKEET_HOP_LENGTH
+    buckets = []
+    for seconds in audio_bucket_seconds():
+        frames = seconds * frames_per_second
+        if frames < max_frames:
+            buckets.append(frames)
+    buckets.append(max_frames)
+    return buckets
+
+
 def normalize_audio_samples(samples: np.ndarray) -> np.ndarray:
     if samples.dtype == np.uint8:
         return ((samples.astype(np.float32) - 128.0) / 128.0).astype(np.float32)

python/cactus/transpile/tdt_runtime.py

@@ -13,6 +13,7 @@
 import torch.nn.functional as F
 
 from cactus.transpile.component_pipeline import ComponentModuleSpec
+from cactus.transpile.audio_preprocess import audio_bucket_frames
 from cactus.transpile.audio_preprocess import prepare_native_parakeet_audio_features
 from cactus.transpile.model_profiles import add_tensor_aliases
 from cactus.transpile.model_profiles import PARAKEET_TDT_PROFILE
@@ -871,16 +872,28 @@ def build_parakeet_tdt_component_specs(
         "adapter_family": "parakeet_tdt",
     }
 
+    full_frames = int(input_features.shape[1])
+    encoder_specs: list[ComponentModuleSpec] = []
+    for frames in audio_bucket_frames(full_frames):
+        component = "audio_encoder" if frames == full_frames else f"audio_encoder_{frames}"
+        encoder_specs.append(
+            ComponentModuleSpec(
+                component=component,
+                module=model.encoder,
+                example_inputs=(input_features[:, :frames, :].contiguous(),),
+                input_keys=("input_features",),
+                output_keys=("encoder_hidden_states",),
+                graph_meta={**common_graph_meta, "component": component},
+                metadata={
+                    "family": "parakeet_tdt",
+                    "task": "tdt_transcription",
+                    "audio_frames": str(frames),
+                },
+            )
+        )
+
     return [
-        ComponentModuleSpec(
-            component="audio_encoder",
-            module=model.encoder,
-            example_inputs=(input_features,),
-            input_keys=("input_features",),
-            output_keys=("encoder_hidden_states",),
-            graph_meta={**common_graph_meta, "component": "audio_encoder"},
-            metadata={"family": "parakeet_tdt", "task": "tdt_transcription"},
-        ),
+        *encoder_specs,
         ComponentModuleSpec(
             component="decoder",
             module=model.decoder_step,

python/pyproject.toml

@@ -73,7 +73,7 @@ exclude = ["cactus.code*"]
 
 [tool.setuptools.package-data]
 "cactus.bindings" = ["lib/*.dylib", "lib/*.so"]
-"cactus" = ["bin/run", "bin/transcribe", "assets/*", "code/**/*"]
+"cactus" = ["bin/run", "bin/transcribe", "assets/*", "code/**/*", "py.typed"]
 "cactus.convert" = ["assets/*/*"]
 
 [tool.pytest.ini_options]