perf(k2_horizon_uno): prefill prompts in 64-row base-profile chunks

families/k2_horizon_uno/model.py

@@ -138,24 +138,33 @@ def _with_width(network, rows, width: int, keepalive: list[np.ndarray]):
 
 LORA_PROFILE = 0
 BASE_PROFILE = 1
+# Prompt rows per base-profile call; decode blocks stay at ``max_block_size``.
+PROMPT_CHUNK_ROWS = 64
 
 
-def add_lora_and_base_profiles(builder, builder_config, row_inputs, max_rows: int) -> None:
+def add_lora_and_base_profiles(
+    builder, builder_config, row_inputs, max_rows: int, base_max_rows: int | None = None
+) -> None:
     """Add the two optimization profiles of a ``1..max_rows`` row engine.
 
     ``LORA_PROFILE`` takes one ``lora_mask`` value per row and runs the LoRA
     branch. ``BASE_PROFILE`` takes an empty mask, for prompt and verify calls
     whose mask is all 0, so the LoRA GEMMs get no rows. The native runtime
-    keeps one execution context per profile.
+    keeps one execution context per profile. ``base_max_rows`` (default
+    ``max_rows``) lets prompt chunks on the base profile exceed a decode block;
+    its optimum stays ``max_rows``, the size of the decode-time base calls.
     """
 
-    for mask_rows in ((1, max_rows, max_rows), (0, 0, 0)):
+    base_max_rows = max_rows if base_max_rows is None else base_max_rows
+    if base_max_rows < max_rows:
+        raise ValueError("K2-Horizon-Uno base profile rows must cover a decode block")
+    for mask_rows, high in (((1, max_rows, max_rows), max_rows), ((0, 0, 0), base_max_rows)):
         profile = builder.create_optimization_profile()
         for name in row_inputs:
             # TensorRT 11.1 mutates the profile in place and returns ``None`` from
             # this Python binding.  Registration below is the authoritative
             # validation boundary.
-            profile.set_shape(name, (1,), (max_rows,), (max_rows,))
+            profile.set_shape(name, (1,), (max_rows,), (high,))
         profile.set_shape("lora_mask", *((rows,) for rows in mask_rows))
         if builder_config.add_optimization_profile(profile) < 0:
             raise RuntimeError("TensorRT rejected the K2-Horizon-Uno optimization profile")
@@ -486,7 +495,11 @@ def build_engine(
         raise RuntimeError("TensorRT failed to create K2-Horizon-Uno engine inputs")
 
     add_lora_and_base_profiles(
-        builder, builder_config, ("token_id", "position_id"), cfg.max_block_size
+        builder,
+        builder_config,
+        ("token_id", "position_id"),
+        cfg.max_block_size,
+        base_max_rows=max(cfg.max_block_size, min(PROMPT_CHUNK_ROWS, max_cache_length)),
     )
     lora_rows = _lora_rows(network, token_id, lora_mask)
 

families/k2_horizon_uno/runtime/kv_cache.cpp

@@ -15,7 +15,7 @@ namespace trtmc {
 namespace {
 
 constexpr int32_t kHeadDim = 128;
-constexpr int32_t kMaximumBlockLength = 8;
+constexpr int32_t kMaximumCallRows = K2HorizonUnoKvCache::kMaximumCallRows;
 constexpr int32_t kNumKvHeads = 8;
 constexpr int32_t kNumLayers = 36;
 
@@ -83,8 +83,8 @@ K2HorizonUnoCacheCursor::K2HorizonUnoCacheCursor(int32_t capacity) : capacity_(c
 }
 
 void K2HorizonUnoCacheCursor::validate_block(int32_t sequence_length) const {
-    if (sequence_length <= 0 || sequence_length > kMaximumBlockLength) {
-        throw std::invalid_argument("K2-Horizon-Uno block length must be in [1, 8]");
+    if (sequence_length <= 0 || sequence_length > kMaximumCallRows) {
+        throw std::invalid_argument("K2-Horizon-Uno call rows must be in [1, 64]");
     }
     if (position_ > capacity_ - sequence_length)
         throw std::runtime_error("K2-Horizon-Uno sequence exceeds fixed KV capacity");

families/k2_horizon_uno/runtime/kv_cache.h

@@ -50,6 +50,9 @@ class K2HorizonUnoCacheCursor {
 // temporary draft/verify suffixes, which later enqueues overwrite.
 class K2HorizonUnoKvCache {
   public:
+    // Rows of one call: decode blocks use at most 8, prompt chunks up to 64.
+    static constexpr int32_t kMaximumCallRows = 64;
+
     K2HorizonUnoKvCache(int32_t max_length, cudaStream_t stream, K2HorizonUnoKvCacheNames names);
 
     void reset();

families/k2_horizon_uno/runtime/pipeline.cpp

@@ -56,6 +56,13 @@ K2HorizonUnoTextGenerationPipeline::K2HorizonUnoTextGenerationPipeline(
     : lora_decoder_(std::move(lora_decoder)), base_decoder_(std::move(base_decoder)),
       cache_(std::move(cache)), tokenizer_(std::move(tokenizer)) {
     validate_pipeline_components(lora_decoder_, base_decoder_, cache_);
+    const auto maximum = base_decoder_->input_profile_shape(
+        "token_id", base_decoder_->profile_idx(), ProfileShapeSelector::kMax);
+    if (maximum.size() != 1 || maximum.front() < 1)
+        throw std::runtime_error("K2HorizonUnoTextGenerationPipeline: invalid base token profile");
+    // A larger base profile is fine, but one cache call takes at most kMaximumCallRows rows.
+    prompt_rows_ = static_cast<int32_t>(
+        std::min<int64_t>(maximum.front(), K2HorizonUnoKvCache::kMaximumCallRows));
     // The engine has no data-dependent shapes or host reads, so a call is a
     // fixed kernel sequence for its row count and a context can replay it as
     // a CUDA graph; a context keeps one graph and recaptures when its row
@@ -236,19 +243,30 @@ void K2HorizonUnoTextGenerationPipeline::reset_generation_context() {
 
 void K2HorizonUnoTextGenerationPipeline::run_prefill(const std::vector<int32_t>& input_ids,
                                                      std::vector<float>& logits) {
-    const std::size_t chunk = kMaximumBlockLength;
+    // Every prompt row but the last only fills the KV cache, in chunks of up to
+    // the base profile's maximum rows.
+    const std::size_t chunk = static_cast<std::size_t>(prompt_rows_);
     std::size_t offset = 0;
     while (input_ids.size() - offset > 1) {
         const std::size_t rows = std::min(chunk, input_ids.size() - offset - 1);
-        const auto begin = input_ids.begin() + static_cast<std::ptrdiff_t>(offset);
-        const std::vector<int32_t> tokens(begin, begin + static_cast<std::ptrdiff_t>(rows));
-        run_block(tokens, std::vector<float>(rows, 0.0F), logits);
+        run_prompt_rows(input_ids.data() + offset, static_cast<int32_t>(rows));
         offset += rows;
     }
     // Keep the final row isolated because seed selection consumes logits row zero.
     run_block({input_ids[offset]}, {0.0F}, logits);
 }
 
+void K2HorizonUnoTextGenerationPipeline::run_prompt_rows(const int32_t* token_ids, int32_t rows) {
+    TensorMap inputs;
+    inputs["token_id"] = Tensor{const_cast<int32_t*>(token_ids), {rows}, DType::kInt32};
+    float empty_mask = 0.0F;
+    inputs["lora_mask"] = Tensor{&empty_mask, {int64_t{0}}, DType::kFloat32};
+    cache_->prepare_block(inputs, rows);
+    base_decoder_->forward_async(inputs);
+    base_decoder_->sync();
+    cache_->advance(rows);
+}
+
 void K2HorizonUnoTextGenerationPipeline::run_block(const std::vector<int32_t>& token_ids,
                                                    const std::vector<float>& lora_mask,
                                                    std::vector<float>& logits) {

families/k2_horizon_uno/runtime/pipeline.h

@@ -58,13 +58,17 @@ class K2HorizonUnoTextGenerationPipeline final : public ITextGeneration {
 
     void reset_generation_context();
     void run_prefill(const std::vector<int32_t>& input_ids, std::vector<float>& logits);
+    // Run prompt rows on the base context without reading their logits.
+    void run_prompt_rows(const int32_t* token_ids, int32_t rows);
     void run_block(const std::vector<int32_t>& token_ids, const std::vector<float>& lora_mask,
                    std::vector<float>& logits);
     std::vector<int32_t> effective_eos_ids(const TextGenerationConfig& cfg) const;
     void log_decode_receipt(const DecodeStats& stats) const;
 
     std::unique_ptr<ITrtModule> lora_decoder_;
     std::unique_ptr<ITrtModule> base_decoder_;
+    // Prompt rows per base-context call: the base profile maximum.
+    int32_t prompt_rows_{1};
     std::unique_ptr<K2HorizonUnoKvCache> cache_;
     std::shared_ptr<ITokenizer> tokenizer_;
     double last_setup_ms_{0.0};

families/k2_horizon_uno/runtime/plugin.cpp

@@ -30,6 +30,8 @@ namespace trtmc::k2_horizon_uno {
 namespace {
 
 constexpr std::int32_t kMaximumBlockLength = 8;
+// Prompt chunks on the base profile; bundles built before them stop at 8.
+constexpr std::int32_t kMaximumPromptRows = 64;
 constexpr std::int32_t kVocabSize = 250624;
 constexpr std::int32_t kNumLayers = 36;
 constexpr std::int32_t kMaxPositionEmbeddings = 524288;
@@ -118,15 +120,20 @@ K2HorizonUnoKvCacheNames make_cache_names() {
 
 void require_dynamic_vector_input(const ITrtModule& module, const std::string& name, DType dtype,
                                   std::int32_t profile) {
+    const auto mismatch = [&name] {
+        return std::runtime_error("K2-Horizon-Uno dynamic input contract mismatch for '" + name +
+                                  "'");
+    };
     if (!module.has_input(name) || module.tensor_dtype(name) != dtype ||
         !module.input_is_dynamic(name) ||
         module.input_profile_shape(name, profile, ProfileShapeSelector::kMin) !=
-            std::vector<std::int64_t>{1} ||
-        module.input_profile_shape(name, profile, ProfileShapeSelector::kMax) !=
-            std::vector<std::int64_t>{kMaximumBlockLength}) {
-        throw std::runtime_error("K2-Horizon-Uno dynamic input contract mismatch for '" + name +
-                                 "'");
-    }
+            std::vector<std::int64_t>{1})
+        throw mismatch();
+    // Decode blocks need 8 rows; the base profile may take longer prompt chunks.
+    const auto maximum = module.input_profile_shape(name, profile, ProfileShapeSelector::kMax);
+    const std::int64_t limit = profile == kBaseProfile ? kMaximumPromptRows : kMaximumBlockLength;
+    if (maximum.size() != 1 || maximum.front() < kMaximumBlockLength || maximum.front() > limit)
+        throw mismatch();
     const auto optimum = module.input_profile_shape(name, profile, ProfileShapeSelector::kOpt);
     if (optimum.size() != 1 || optimum.front() < 1 || optimum.front() > kMaximumBlockLength) {
         throw std::runtime_error("K2-Horizon-Uno input profile optimum is invalid for '" + name +

families/k2_horizon_uno/tests/cpp/test_k2_horizon_uno_algorithm.cpp

@@ -227,7 +227,11 @@ void test_cache_cursor_rollback_contract() {
     check(rejects([&] { cursor.rollback(4); }), "rollback cannot move frontier forward");
     cursor.reset();
     check(cursor.position() == 0, "cursor reset clears logical state");
-    check(rejects([&] { cursor.advance(9); }), "cursor rejects blocks above engine maximum");
+    // Decode blocks use at most 8 rows; prompt chunks of the base profile up to 64.
+    trtmc::K2HorizonUnoCacheCursor prompt(128);
+    check(rejects([&] { prompt.advance(65); }), "cursor rejects calls above engine maximum");
+    prompt.advance(64);
+    check(prompt.position() == 64, "prompt chunk advances cursor by up to 64 rows");
 
     // Three prompt rows plus five completion tokens fill capacity exactly.
     // The first completion seed remains uncached: draft reaches 7, then

families/k2_horizon_uno/tests/test_lora_profiles.py

@@ -26,14 +26,16 @@ def _from_bits(bits: np.ndarray) -> torch.Tensor:
     return torch.from_numpy(bits.view(np.int16)).view(torch.bfloat16).float()
 
 
-def _build_projection_engine(table, weight, lora_a, lora_b):
+def _build_projection_engine(table, weight, lora_a, lora_b, base_max_rows=None):
     logger = trt.Logger(trt.Logger.WARNING)
     builder = trt.Builder(logger)
     network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
     config = builder.create_builder_config()
     token_id = network.add_input("token_id", trt.int32, (-1,))
     lora_mask = network.add_input("lora_mask", trt.float32, (-1,))
-    model.add_lora_and_base_profiles(builder, config, ("token_id",), _ROWS)
+    model.add_lora_and_base_profiles(
+        builder, config, ("token_id",), _ROWS, base_max_rows=base_max_rows
+    )
     keepalive: list[np.ndarray] = []
     embedding = model._work_constant(
         network, table.shape, table, work_dtype=trt.bfloat16, constant_keepalive=keepalive
@@ -111,3 +113,35 @@ def test_base_profile_skips_lora_and_lora_profile_applies_the_mask():
         torch.testing.assert_close(skipped, base[:rows], rtol=2e-2, atol=2e-2)
         if rows == 8:
             assert torch.equal(skipped, zero_mask)
+
+
+@pytest.mark.gpu
+@pytest.mark.trt
+def test_base_profile_takes_prompt_chunks_longer_than_a_decode_block():
+    """The base profile can run more rows than a decode block; the LoRA profile cannot."""
+    rng = np.random.default_rng(1)
+    table = _bf16_bits(rng.standard_normal((_VOCAB, _IN)))
+    weight = _bf16_bits(rng.standard_normal((_IN, _OUT)) * 0.1)
+    lora_a = _bf16_bits(rng.standard_normal((_IN, _RANK)) * 0.1)
+    lora_b = _bf16_bits(rng.standard_normal((_RANK, _OUT)) * 0.1)
+    engine = _build_projection_engine(table, weight, lora_a, lora_b, base_max_rows=4 * _ROWS)
+    lora_shapes = engine.get_tensor_profile_shape("token_id", model.LORA_PROFILE)
+    base_shapes = engine.get_tensor_profile_shape("token_id", model.BASE_PROFILE)
+    assert [tuple(s) for s in lora_shapes] == [(1,), (_ROWS,), (_ROWS,)]
+    assert [tuple(s) for s in base_shapes] == [(1,), (_ROWS,), (4 * _ROWS,)]
+
+    tokens = rng.integers(0, _VOCAB, 4 * _ROWS - 3).tolist()
+    expected = _from_bits(table)[tokens] @ _from_bits(weight)
+    got = _run(engine, model.BASE_PROFILE, tokens, [])
+    torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2)
+
+
+@pytest.mark.gpu
+@pytest.mark.trt
+def test_base_profile_must_cover_a_decode_block():
+    """A base profile smaller than a decode block is rejected."""
+    builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
+    with pytest.raises(ValueError, match="decode block"):
+        model.add_lora_and_base_profiles(
+            builder, builder.create_builder_config(), ("token_id",), _ROWS, base_max_rows=_ROWS - 1
+        )