From 80814651db5b20b7fca29335136c231ce12c38fa Mon Sep 17 00:00:00 2001 From: chaofengw Date: Tue, 10 Nov 2026 17:09:20 +0000 Subject: [PATCH 1/2] perf(k2_horizon_uno): prefill prompts in 64-row base-profile chunks The base optimization profile (empty LoRA mask) now accepts up to 64 rows while its optimum and the LoRA profile stay at the 8-row decode block, so the runtime fills the KV cache with 64-row prompt chunks instead of 8-row blocks. Bundles built before this change keep an 8-row base profile and the runtime still accepts them. Signed-off-by: chaofengw --- families/k2_horizon_uno/model.py | 23 ++++++++--- families/k2_horizon_uno/runtime/kv_cache.cpp | 7 ++-- families/k2_horizon_uno/runtime/pipeline.cpp | 24 ++++++++++-- families/k2_horizon_uno/runtime/pipeline.h | 4 ++ families/k2_horizon_uno/runtime/plugin.cpp | 19 +++++++--- .../cpp/test_k2_horizon_uno_algorithm.cpp | 6 ++- .../tests/test_lora_profiles.py | 38 ++++++++++++++++++- 7 files changed, 100 insertions(+), 21 deletions(-) diff --git a/families/k2_horizon_uno/model.py b/families/k2_horizon_uno/model.py index a04511cc99..81a7d5f599 100644 --- a/families/k2_horizon_uno/model.py +++ b/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) diff --git a/families/k2_horizon_uno/runtime/kv_cache.cpp b/families/k2_horizon_uno/runtime/kv_cache.cpp index 130e44c4f8..3723d2d9e8 100644 --- a/families/k2_horizon_uno/runtime/kv_cache.cpp +++ b/families/k2_horizon_uno/runtime/kv_cache.cpp @@ -15,7 +15,8 @@ namespace trtmc { namespace { constexpr int32_t kHeadDim = 128; -constexpr int32_t kMaximumBlockLength = 8; +// Rows of one call: decode blocks use at most 8, prompt chunks up to 64. +constexpr int32_t kMaximumCallRows = 64; constexpr int32_t kNumKvHeads = 8; constexpr int32_t kNumLayers = 36; @@ -83,8 +84,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"); diff --git a/families/k2_horizon_uno/runtime/pipeline.cpp b/families/k2_horizon_uno/runtime/pipeline.cpp index c846e11f52..ef43b0ba79 100644 --- a/families/k2_horizon_uno/runtime/pipeline.cpp +++ b/families/k2_horizon_uno/runtime/pipeline.cpp @@ -56,6 +56,11 @@ 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"); + prompt_rows_ = static_cast(maximum.front()); // 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 +241,30 @@ void K2HorizonUnoTextGenerationPipeline::reset_generation_context() { void K2HorizonUnoTextGenerationPipeline::run_prefill(const std::vector& input_ids, std::vector& 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(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(offset); - const std::vector tokens(begin, begin + static_cast(rows)); - run_block(tokens, std::vector(rows, 0.0F), logits); + run_prompt_rows(input_ids.data() + offset, static_cast(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(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& token_ids, const std::vector& lora_mask, std::vector& logits) { diff --git a/families/k2_horizon_uno/runtime/pipeline.h b/families/k2_horizon_uno/runtime/pipeline.h index 3452367abf..8adae373ef 100644 --- a/families/k2_horizon_uno/runtime/pipeline.h +++ b/families/k2_horizon_uno/runtime/pipeline.h @@ -58,6 +58,8 @@ class K2HorizonUnoTextGenerationPipeline final : public ITextGeneration { void reset_generation_context(); void run_prefill(const std::vector& input_ids, std::vector& 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& token_ids, const std::vector& lora_mask, std::vector& logits); std::vector effective_eos_ids(const TextGenerationConfig& cfg) const; @@ -65,6 +67,8 @@ class K2HorizonUnoTextGenerationPipeline final : public ITextGeneration { std::unique_ptr lora_decoder_; std::unique_ptr base_decoder_; + // Prompt rows per base-context call: the base profile maximum. + int32_t prompt_rows_{1}; std::unique_ptr cache_; std::shared_ptr tokenizer_; double last_setup_ms_{0.0}; diff --git a/families/k2_horizon_uno/runtime/plugin.cpp b/families/k2_horizon_uno/runtime/plugin.cpp index d8c308d325..7f4c8b1a37 100644 --- a/families/k2_horizon_uno/runtime/plugin.cpp +++ b/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{1} || - module.input_profile_shape(name, profile, ProfileShapeSelector::kMax) != - std::vector{kMaximumBlockLength}) { - throw std::runtime_error("K2-Horizon-Uno dynamic input contract mismatch for '" + name + - "'"); - } + std::vector{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 + diff --git a/families/k2_horizon_uno/tests/cpp/test_k2_horizon_uno_algorithm.cpp b/families/k2_horizon_uno/tests/cpp/test_k2_horizon_uno_algorithm.cpp index 7635d7a1e9..173180cc07 100644 --- a/families/k2_horizon_uno/tests/cpp/test_k2_horizon_uno_algorithm.cpp +++ b/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 diff --git a/families/k2_horizon_uno/tests/test_lora_profiles.py b/families/k2_horizon_uno/tests/test_lora_profiles.py index 14c4c2ec60..c10f10b997 100644 --- a/families/k2_horizon_uno/tests/test_lora_profiles.py +++ b/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 + ) From 79e4f67dcb72f4e757929b232e8758acd86b1c87 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Wed, 11 Nov 2026 06:50:41 +0000 Subject: [PATCH 2/2] fix(k2_horizon_uno): cap prompt chunks at the cache's call-row limit A pipeline built without the plugin's profile validation could take a base profile larger than 64 rows and fail in prepare_block on a long prompt. Cap the prompt chunk at K2HorizonUnoKvCache::kMaximumCallRows, now the single definition of that limit. Signed-off-by: chaofengw --- families/k2_horizon_uno/runtime/kv_cache.cpp | 3 +-- families/k2_horizon_uno/runtime/kv_cache.h | 3 +++ families/k2_horizon_uno/runtime/pipeline.cpp | 4 +++- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/families/k2_horizon_uno/runtime/kv_cache.cpp b/families/k2_horizon_uno/runtime/kv_cache.cpp index 3723d2d9e8..c687fb2f8b 100644 --- a/families/k2_horizon_uno/runtime/kv_cache.cpp +++ b/families/k2_horizon_uno/runtime/kv_cache.cpp @@ -15,8 +15,7 @@ namespace trtmc { namespace { constexpr int32_t kHeadDim = 128; -// Rows of one call: decode blocks use at most 8, prompt chunks up to 64. -constexpr int32_t kMaximumCallRows = 64; +constexpr int32_t kMaximumCallRows = K2HorizonUnoKvCache::kMaximumCallRows; constexpr int32_t kNumKvHeads = 8; constexpr int32_t kNumLayers = 36; diff --git a/families/k2_horizon_uno/runtime/kv_cache.h b/families/k2_horizon_uno/runtime/kv_cache.h index 1c5249e309..2880a064a8 100644 --- a/families/k2_horizon_uno/runtime/kv_cache.h +++ b/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(); diff --git a/families/k2_horizon_uno/runtime/pipeline.cpp b/families/k2_horizon_uno/runtime/pipeline.cpp index ef43b0ba79..8e90237de6 100644 --- a/families/k2_horizon_uno/runtime/pipeline.cpp +++ b/families/k2_horizon_uno/runtime/pipeline.cpp @@ -60,7 +60,9 @@ K2HorizonUnoTextGenerationPipeline::K2HorizonUnoTextGenerationPipeline( "token_id", base_decoder_->profile_idx(), ProfileShapeSelector::kMax); if (maximum.size() != 1 || maximum.front() < 1) throw std::runtime_error("K2HorizonUnoTextGenerationPipeline: invalid base token profile"); - prompt_rows_ = static_cast(maximum.front()); + // A larger base profile is fine, but one cache call takes at most kMaximumCallRows rows. + prompt_rows_ = static_cast( + std::min(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