perf(voxcpm2): size the step graphs by the request, not by max_length

include/engine/framework/runtime/kv_cache.h

@@ -2,6 +2,7 @@
 
 #include "engine/framework/core/module.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <string>
 #include <vector>
@@ -29,6 +30,27 @@ inline int64_t ring_stored_position(
     return row < pinned_prefix_steps ? row : current_end - valid_steps + row;
 }
 
+// Bucketed capacity for a static-cache AR loop: prompt rows plus a small
+// generated window to start, doubling buckets up to prompt + max_tokens. A
+// family that sizes its step graph from this number builds (and attends over)
+// the cache the request can fill, not the model's full cache length.
+inline int64_t bucketed_cache_steps(
+    int64_t prompt_steps,
+    int64_t max_tokens,
+    int64_t initial_generated_steps = 128,
+    int64_t minimum_bucket_steps = 128) {
+    const int64_t maximum = prompt_steps + max_tokens;
+    const int64_t required = prompt_steps + std::min(max_tokens, initial_generated_steps);
+    int64_t bucket = minimum_bucket_steps;
+    while (bucket < required && bucket <= maximum / 2) {
+        bucket *= 2;
+    }
+    if (bucket < required) {
+        bucket = required;
+    }
+    return std::min(bucket, maximum);
+}
+
 struct TransformerKVCacheOptions {
     bool allow_f16_storage = false;
     bool allow_bf16_storage = false;

include/engine/models/voxcpm2/minicpm.h

@@ -167,6 +167,8 @@ class VoxCPM2MiniCPMStepRuntime final {
   void reset();
   void import_state(const engine::runtime::TransformerKVState &state);
   engine::runtime::TransformerKVState export_state() const;
+  int64_t cache_steps() const;
+  int64_t valid_steps() const;
   VoxCPM2MiniCPMStepOutput run_step(const std::vector<float> &embedding);
   void release_runtime_memory();
 

src/models/higgs_audio_tts/generator.cpp

@@ -2,6 +2,7 @@
 
 #include "engine/framework/debug/profiler.h"
 #include "engine/framework/debug/trace.h"
+#include "engine/framework/runtime/kv_cache.h"
 #include "engine/framework/runtime/options.h"
 #include "engine/framework/sampling/torch_random.h"
 #include "engine/models/higgs_audio_tts/codebooks.h"
@@ -20,22 +21,6 @@ namespace {
 
 using Clock = std::chrono::steady_clock;
 
-constexpr int64_t kInitialGeneratedCacheSteps = 128;
-constexpr int64_t kMinimumCacheBucketSteps = 128;
-
-int64_t bucketed_initial_cache_steps(int64_t prompt_steps, int64_t max_tokens) {
-    const int64_t maximum = prompt_steps + max_tokens;
-    const int64_t required = prompt_steps + std::min(max_tokens, kInitialGeneratedCacheSteps);
-    int64_t bucket = kMinimumCacheBucketSteps;
-    while (bucket < required && bucket <= maximum / 2) {
-        bucket *= 2;
-    }
-    if (bucket < required) {
-        bucket = required;
-    }
-    return std::min(bucket, maximum);
-}
-
 void validate_generation_options(const HiggsGenerationOptions & options) {
     if (options.max_tokens <= 0) {
         throw std::runtime_error("Higgs TTS max_tokens must be positive after default resolution");
@@ -391,7 +376,8 @@ HiggsGenerationResult HiggsGenerator::generate(
     engine::debug::trace_log_scalar("higgs_audio_tts.generator.reference_prefix_cache_hit", reference_cache_hit);
     engine::debug::trace_log_scalar("higgs_audio_tts.generator.reference_prefix_steps", prepared.prefix_steps);
     const int64_t max_cache_steps = prompt_steps + request.options.max_tokens;
-    const int64_t initial_cache_steps = bucketed_initial_cache_steps(prompt_steps, request.options.max_tokens);
+    const int64_t initial_cache_steps =
+        engine::runtime::bucketed_cache_steps(prompt_steps, request.options.max_tokens);
     const bool cache_rebuild =
         ar_kv_cache_ == nullptr || !ar_kv_cache_->can_run(*ar_, initial_cache_steps) ||
         ar_kv_cache_->cache_steps() != initial_cache_steps;

src/models/voxcpm2/generator.cpp

@@ -12,6 +12,7 @@
 #include "engine/framework/modules/weight_binding.h"
 #include "engine/framework/runtime/cache_slots.h"
 #include "engine/framework/runtime/errors.h"
+#include "engine/framework/runtime/kv_cache.h"
 #include "engine/framework/sampling/torch_random.h"
 #include "engine/models/voxcpm2/assets.h"
 #include "engine/models/voxcpm2/minicpm.h"
@@ -1210,12 +1211,7 @@ class VoxCPM2FeatureGeneratorRuntime::Impl {
         text_embedding_(weights_, config.text_embedding_graph_context_bytes,
                         config.mem_saver),
         prefill_(weights_, config.lm_step_graph_context_bytes),
-        base_lm_(weights_, VoxCPM2MiniCPMKind::BaseLM,
-                 assets_->config.max_length,
-                 config.lm_step_graph_context_bytes),
-        residual_lm_(weights_, VoxCPM2MiniCPMKind::ResidualLM,
-                     assets_->config.max_length,
-                     config.lm_step_graph_context_bytes),
+        lm_step_graph_context_bytes_(config.lm_step_graph_context_bytes),
         projection_(weights_, config.projection_graph_context_bytes,
                     config.mem_saver),
         cfm_(weights_, config.dit_graph_context_bytes, config.mem_saver),
@@ -1280,8 +1276,36 @@ class VoxCPM2FeatureGeneratorRuntime::Impl {
 
   void release_runtime_memory() {
     prefill_.release_runtime_memory();
-    base_lm_.release_runtime_memory();
-    residual_lm_.release_runtime_memory();
+    if (base_lm_) {
+      base_lm_->release_runtime_memory();
+    }
+    if (residual_lm_) {
+      residual_lm_->release_runtime_memory();
+    }
+  }
+
+  // The step graphs are sized by what the request can fill, not by the model's
+  // cache length: cache_steps_ is a graph dimension, so an 8192-row cache means
+  // every generated patch attends over 8192 rows and 36 layers x 8192 x 256
+  // f32 keys and values stay resident (1.21 GiB on the shipped package) even for
+  // a three-second phrase. Rebuild only when the bucket changes.
+  void replace_step_runtimes(int64_t cache_steps, bool preserve_state) {
+    const auto base_state =
+        preserve_state && base_lm_ ? base_lm_->export_state()
+                                   : engine::runtime::TransformerKVState{};
+    const auto residual_state =
+        preserve_state && residual_lm_ ? residual_lm_->export_state()
+                                       : engine::runtime::TransformerKVState{};
+    base_lm_ = std::make_unique<VoxCPM2MiniCPMStepRuntime>(
+        weights_, VoxCPM2MiniCPMKind::BaseLM, cache_steps,
+        lm_step_graph_context_bytes_);
+    residual_lm_ = std::make_unique<VoxCPM2MiniCPMStepRuntime>(
+        weights_, VoxCPM2MiniCPMKind::ResidualLM, cache_steps,
+        lm_step_graph_context_bytes_);
+    if (preserve_state) {
+      base_lm_->import_state(base_state);
+      residual_lm_->import_state(residual_state);
+    }
   }
 
 private:
@@ -1507,8 +1531,15 @@ class VoxCPM2FeatureGeneratorRuntime::Impl {
           + std::to_string(max_tokens) + " requested tokens, limit "
           + std::to_string(config.max_length) + "); shorten the input text");
     }
-    base_lm_.reset();
-    residual_lm_.reset();
+    const int64_t prefill_steps = static_cast<int64_t>(prefill.rows.size());
+    const int64_t max_cache_steps = prefill_steps + max_tokens;
+    const int64_t cache_steps =
+        engine::runtime::bucketed_cache_steps(prefill_steps, max_tokens);
+    if (base_lm_ == nullptr || base_lm_->cache_steps() != cache_steps) {
+      replace_step_runtimes(cache_steps, false);
+    }
+    base_lm_->reset();
+    residual_lm_->reset();
 
     std::vector<float> zero_hidden(static_cast<size_t>(hidden_size), 0.0F);
     std::vector<float> zero_patch(static_cast<size_t>(patch_elems), 0.0F);
@@ -1548,8 +1579,8 @@ class VoxCPM2FeatureGeneratorRuntime::Impl {
       prefill_input.audio_mask.push_back(row.audio_mask ? 1.0F : 0.0F);
     }
     const auto prefill_output = prefill_.run(prefill_input);
-    base_lm_.import_state(prefill_output.base_state);
-    residual_lm_.import_state(prefill_output.residual_state);
+    base_lm_->import_state(prefill_output.base_state);
+    residual_lm_->import_state(prefill_output.residual_state);
     std::vector<float> lm_hidden = prefill_output.lm_hidden;
     std::vector<float> residual_hidden = prefill_output.residual_hidden;
 
@@ -1621,12 +1652,29 @@ class VoxCPM2FeatureGeneratorRuntime::Impl {
       }
 
       const auto curr_embed = local_encoder_.encode_patch(patch);
-      const auto next_lm = base_lm_.run_step(curr_embed).hidden;
+      if (base_lm_->valid_steps() >= base_lm_->cache_steps()) {
+        // The bucket filled up: rebuild both step graphs at the next bucket and
+        // carry the state over. max_cache_steps is the request's own bound, so
+        // the last bucket is exactly what the request can fill.
+        const int64_t grown = std::min<int64_t>(
+            max_cache_steps,
+            std::max<int64_t>(base_lm_->cache_steps() * 2,
+                              base_lm_->valid_steps() + 1));
+        if (grown <= base_lm_->cache_steps()) {
+          throw engine::runtime::CapacityError(
+              "VoxCPM2 generation exceeds the model cache length ("
+              + std::to_string(base_lm_->valid_steps()) + " cache rows filled, "
+              + std::to_string(max_cache_steps) + " available, limit "
+              + std::to_string(config.max_length) + "); raise max_tokens");
+        }
+        replace_step_runtimes(grown, true);
+      }
+      const auto next_lm = base_lm_->run_step(curr_embed).hidden;
       const auto next_projected =
           projection_.run(next_lm, residual_hidden, curr_embed);
       lm_hidden = next_projected.fsq_hidden;
       residual_hidden =
-          residual_lm_.run_step(next_projected.residual_input).hidden;
+          residual_lm_->run_step(next_projected.residual_input).hidden;
     }
     return result;
   }
@@ -1636,8 +1684,9 @@ class VoxCPM2FeatureGeneratorRuntime::Impl {
   VoxCPM2TextTokenizer tokenizer_;
   VoxCPM2TextEmbeddingRuntime text_embedding_;
   VoxCPM2PromptPrefillRuntime prefill_;
-  VoxCPM2MiniCPMStepRuntime base_lm_;
-  VoxCPM2MiniCPMStepRuntime residual_lm_;
+  size_t lm_step_graph_context_bytes_ = 0;
+  std::unique_ptr<VoxCPM2MiniCPMStepRuntime> base_lm_;
+  std::unique_ptr<VoxCPM2MiniCPMStepRuntime> residual_lm_;
   VoxCPM2StepProjectionRuntime projection_;
   VoxCPM2CFMRuntime cfm_;
   VoxCPM2LocalEncoderRuntime local_encoder_;

src/models/voxcpm2/minicpm.cpp

@@ -1016,6 +1016,9 @@ class VoxCPM2MiniCPMStepRuntime::Impl {
     step_cache_.import_state(state);
   }
 
+  int64_t cache_steps() const { return cache_steps_; }
+  int64_t valid_steps() const { return step_cache_.valid_steps(); }
+
   void import_state(const engine::runtime::TransformerKVState &state) {
     ensure_graph();
     step_cache_.import_state(state);
@@ -1207,6 +1210,14 @@ VoxCPM2MiniCPMStepRuntime::export_state() const {
   return impl_->export_state();
 }
 
+int64_t VoxCPM2MiniCPMStepRuntime::cache_steps() const {
+  return impl_->cache_steps();
+}
+
+int64_t VoxCPM2MiniCPMStepRuntime::valid_steps() const {
+  return impl_->valid_steps();
+}
+
 VoxCPM2MiniCPMStepOutput
 VoxCPM2MiniCPMStepRuntime::run_step(const std::vector<float> &embedding) {
   return impl_->run_step(embedding);

tests/unittests/test_transformer_kv_ring.cpp

@@ -1,12 +1,14 @@
 // Unit tests for TransformerKVCache ring_mode: slot mapping, wrap advance,
-// chronological export, wrap-placement import, and non-ring regression.
+// chronological export, wrap-placement import, and non-ring regression, plus
+// the bucketed step-graph capacity a static-cache AR loop sizes itself from.
 #include "engine/framework/runtime/kv_cache.h"
 
 #include "engine/framework/core/backend.h"
 
 #include "ggml-backend.h"
 #include "ggml.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <iostream>
 #include <memory>
@@ -265,6 +267,31 @@ void test_non_ring_regression() {
     std::cout << "non-ring regression done\n";
 }
 
+void test_bucketed_cache_steps() {
+    // The shape families size their step graph from: doubling buckets from 128,
+    // clamped to prompt + max_tokens, exact once doubling would overshoot.
+    check(runtime::bucketed_cache_steps(114, 128) == 242, "one generated window is exact");
+    check(runtime::bucketed_cache_steps(100, 200) == 256, "short request rounds to a bucket");
+    check(runtime::bucketed_cache_steps(200, 8000) == 512, "generated window stops at 128 steps");
+    check(runtime::bucketed_cache_steps(8000, 0) == 8000, "long prompt lands on its own length");
+    check(runtime::bucketed_cache_steps(0, 64) == 64, "promptless request clamps to the total");
+    check(runtime::bucketed_cache_steps(0, 0) == 0, "empty request is empty");
+    check(runtime::bucketed_cache_steps(10, 10, 4, 8) == 16, "window and floor knobs are honoured");
+    for (const int64_t prompt : {0, 1, 43, 114, 1000, 8000}) {
+        int64_t previous = 0;
+        for (const int64_t tokens : {0, 1, 64, 128, 129, 1024, 8192}) {
+            const int64_t maximum = prompt + tokens;
+            const int64_t bucket = runtime::bucketed_cache_steps(prompt, tokens);
+            check(bucket <= maximum, "bucket never exceeds the request total");
+            check(bucket >= std::min(prompt + std::min(tokens, int64_t{128}), maximum),
+                  "bucket never undercuts the generated window it needs");
+            check(bucket >= previous, "bucket grows with the request");
+            previous = bucket;
+        }
+    }
+    std::cout << "bucketed_cache_steps done\n";
+}
+
 }  // namespace
 
 int main() {
@@ -274,6 +301,7 @@ int main() {
         test_wrap_round_trip();
         test_wrap_advance();
         test_non_ring_regression();
+        test_bucketed_cache_steps();
     } catch (const std::exception & e) {
         std::cerr << "ERROR: " << e.what() << "\n";
         return 1;