From e12ddc5fa8c81350c133c55ddfb8757869a40955 Mon Sep 17 00:00:00 2001 From: Zhenshan Xie Date: Mon, 9 Nov 2026 10:46:57 -0800 Subject: [PATCH] perf(k2_horizon_uno): skip the LoRA branch for all-0 masks and graph draft calls The engine gets a second optimization profile with an empty lora_mask, so prompt, verify and autoregressive calls (mask all 0) run without the LoRA GEMMs. The native runtime keeps one execution context per profile on one deserialized engine and the same KV cache, picks the context from the mask, and replays draft calls as a CUDA graph. Base-profile calls are bit-identical to the all-0 mask on the old engine and 1.21x to 1.24x faster. Native speculative decode keeps the same tokens and is 1.07x (B200) and 1.02x (B300) faster, prompt time 1.17x to 1.20x, autoregressive mode 1.16x. Bundles must be rebuilt. Addresses part of #2054. Signed-off-by: Zhenshan Xie --- families/k2_horizon_uno/README.md | 40 +++++++ families/k2_horizon_uno/model.py | 105 +++++++++++++--- families/k2_horizon_uno/runtime/algorithm.cpp | 6 + families/k2_horizon_uno/runtime/algorithm.h | 5 + families/k2_horizon_uno/runtime/pipeline.cpp | 41 +++++-- families/k2_horizon_uno/runtime/pipeline.h | 9 +- families/k2_horizon_uno/runtime/plugin.cpp | 67 +++++++---- .../cpp/test_k2_horizon_uno_algorithm.cpp | 6 + .../tests/test_lora_profiles.py | 113 ++++++++++++++++++ 9 files changed, 346 insertions(+), 46 deletions(-) create mode 100644 families/k2_horizon_uno/tests/test_lora_profiles.py diff --git a/families/k2_horizon_uno/README.md b/families/k2_horizon_uno/README.md index 5397c30017..cd118f448d 100644 --- a/families/k2_horizon_uno/README.md +++ b/families/k2_horizon_uno/README.md @@ -10,6 +10,45 @@ decoding: per cycle one draft call (`[seed, 7 noise tokens]`, mask equals greedy autoregressive decoding of the base model. The manifest builds the model in bf16 with `max_sequence_length` 256. Batch size is 1. +## Engine profiles and CUDA graph + +The engine has two optimization profiles over the same weights. Profile 0 +takes one `lora_mask` value per row and runs the LoRA branch; profile 1 takes +an empty `lora_mask` and gives the LoRA GEMMs no rows. The native runtime +deserializes the engine once, keeps one execution context per profile on the +same KV cache, and sends a call to profile 0 only when its mask has a nonzero +value (draft calls). Prompt, verify and autoregressive calls (mask all 0) use +profile 1; their logits and K/V are bit-identical to the same call with an +all-0 mask on profile 0 and on the single-profile engine of earlier versions. +Draft calls on profile 0 run the LoRA rows through other GEMM kernels than +before, so draft logits differ in the last bits (max abs 0.4 to 1.8 on one +test call). This changes only which draft tokens are proposed at near ties; +the committed tokens are the verify argmax and did not change on any measured +prompt (48 of 48 runs; 3150 to 3162 forwards against 3156). Bundles built +before this change have one profile and must be rebuilt. + +The draft context replays a CUDA graph (8 rows in every cycle except the +last). The base context runs without a graph: its row count changes in every +request (prompt chunks of 8, k and 1 rows, then verify blocks), and a recapture +(2 to 13 ms) costs more than a replay saves (about 0.2 ms per call). + +One call, 8 rows, `max_sequence_length` 256, TensorRT 11.4, ms (in brackets: +1024): + +| Call | B200 before | B200 now | B300 before | B300 now | +|---|---|---|---|---| +| Draft (mask `[0, 1, ..., 1]`) | 5.51 (5.64) | 4.73 (4.88), graph | 5.06 (5.12) | 4.91 (5.00), graph | +| Verify or prompt (mask all 0) | 5.51 (5.63) | 4.53 (4.59) | 5.04 (5.12) | 4.16 (4.18) | +| 1 row (AR, last prompt row) | 5.61 | 4.56 | 5.13 | 4.13 | + +Native runtime, 16 chat prompts x 3 rounds, 128 new tokens, greedy, same +tokens as before on every run: speculative decode 141.6 -> 151.6 tokens/s +(B200, 1.07x) and 148.9 -> 151.6 (B300, 1.02x); prompt time per request 26.6 +-> 22.2 ms and 24.2 -> 20.6 ms; autoregressive mode 1.16x on both GPUs. The +end-to-end gain is smaller than the per-call gain because the host side of a +call (copy of the float32 logits and the argmax over 8 x 250624 values on the +CPU, about 7 ms per 8-row call) is longer than the engine time. + ## FP8 quantization study (not shipped) FP8 self-quantization of the base projections was built and measured on @@ -189,6 +228,7 @@ Not shipped: the quality loss is large and speculative decode is slower. the branch is 1.21x to 1.24x faster. Skipping the branch when the mask is all 0 (or a second engine profile without LoRA for prompt and verify) would speed up bf16 decode more than FP8 does. + Addressed: see "Engine profiles and CUDA graph" (item 1 for draft calls). 3. TensorRT 11.4 stream assignment fails for some FP8 graphs of this family (`[MS] Multi stream is disabled because the stream assignment failed`). Without a CUDA graph the multi-stream plans are 1.3x to 1.5x slower than diff --git a/families/k2_horizon_uno/model.py b/families/k2_horizon_uno/model.py index b9e5ba6b9d..a04511cc99 100644 --- a/families/k2_horizon_uno/model.py +++ b/families/k2_horizon_uno/model.py @@ -120,6 +120,74 @@ def _matmul( return layer.get_output(0) +def _shape_length(network, tensor): + shape = network.add_shape(tensor).get_output(0) + return network.add_slice(shape, start=(0,), shape=(1,), stride=(1,)).get_output(0) + + +def _with_width(network, rows, width: int, keepalive: list[np.ndarray]): + """INT64 ``[rows, width]`` shape tensor.""" + + width_tensor = _constant( + network, (1,), np.array([width], dtype=np.int64), keepalive, dtype=np.dtype(np.int64) + ) + layer = network.add_concatenation([rows, width_tensor]) + layer.axis = 0 + return layer.get_output(0) + + +LORA_PROFILE = 0 +BASE_PROFILE = 1 + + +def add_lora_and_base_profiles(builder, builder_config, row_inputs, max_rows: int) -> 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. + """ + + for mask_rows in ((1, max_rows, max_rows), (0, 0, 0)): + 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("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") + + +def _lora_rows(network, token_id, lora_mask): + """Row split of a call, as INT64 ``[1]`` shape tensors ``(rows, remaining)``. + + ``rows`` is the length of ``lora_mask``: the leading token rows that run + the LoRA branch. ``remaining`` is the other token rows, which keep the base + projection. The LoRA profile binds one mask value per token row + (``remaining`` is 0); the base profile binds an empty mask, so the LoRA + GEMMs get no rows. + """ + + rows = _shape_length(network, lora_mask) + remaining = network.add_elementwise( + _shape_length(network, token_id), rows, trt.ElementWiseOperation.SUB + ).get_output(0) + return rows, remaining + + +def _row_slice(network, tensor, start, rows, width: int, keepalive: list[np.ndarray]): + """Rows ``[start, start + rows)`` of a ``[S, width]`` tensor; ``start`` is ``None`` for 0.""" + + layer = network.add_slice(tensor, start=(0, 0), shape=(1, width), stride=(1, 1)) + if start is not None: + layer.set_input(1, _with_width(network, start, 0, keepalive)) + layer.set_input(2, _with_width(network, rows, width, keepalive)) + return layer.get_output(0) + + def _conditional_projection( network, lhs, @@ -127,6 +195,7 @@ def _conditional_projection( lora_a: np.ndarray, lora_b: np.ndarray, lora_mask, + lora_rows, *, lhs_width: int, rhs_width: int, @@ -136,7 +205,9 @@ def _conditional_projection( constant_keepalive: list[np.ndarray], name: str, ): - """Compute ``xW + mask * scale * (xA)B`` for every active token row.""" + """Compute ``xW + mask * scale * (xA)B`` on the mask rows and ``xW`` on the other rows.""" + + rows, remaining = lora_rows base = _matmul( network, @@ -150,7 +221,7 @@ def _conditional_projection( ) low_rank = _matmul( network, - lhs, + _row_slice(network, lhs, None, rows, lhs_width, constant_keepalive), lora_a, lhs_width=lhs_width, rhs_width=rank, @@ -179,16 +250,26 @@ def _conditional_projection( if scaled is None: raise RuntimeError(f"TensorRT failed to scale conditional LoRA for {name}") mask_rows = network.add_shuffle(lora_mask) - mask_rows.reshape_dims = (-1, 1) + # An empty mask cannot use a -1 reshape; give the row count explicitly. + mask_rows.set_input(1, _with_width(network, rows, 1, constant_keepalive)) mask_work = _cast(network, mask_rows.get_output(0), work_dtype) selected = network.add_elementwise( scaled.get_output(0), mask_work, trt.ElementWiseOperation.PROD ) if selected is None: raise RuntimeError(f"TensorRT failed to mask conditional LoRA for {name}") - result = network.add_elementwise(base, selected.get_output(0), trt.ElementWiseOperation.SUM) - if result is None: + lora_part = network.add_elementwise( + _row_slice(network, base, None, rows, rhs_width, constant_keepalive), + selected.get_output(0), + trt.ElementWiseOperation.SUM, + ) + if lora_part is None: raise RuntimeError(f"TensorRT failed to apply conditional LoRA for {name}") + base_part = _row_slice(network, base, rows, remaining, rhs_width, constant_keepalive) + result = network.add_concatenation([lora_part.get_output(0), base_part]) + if result is None: + raise RuntimeError(f"TensorRT failed to join conditional LoRA rows for {name}") + result.axis = 0 result.name = name return result.get_output(0) @@ -404,15 +485,10 @@ def build_engine( ): raise RuntimeError("TensorRT failed to create K2-Horizon-Uno engine inputs") - profile = builder.create_optimization_profile() - optimum = cfg.max_block_size - for name in ("token_id", "position_id", "lora_mask"): - # 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,), (optimum,), (cfg.max_block_size,)) - if builder_config.add_optimization_profile(profile) < 0: - raise RuntimeError("TensorRT rejected the K2-Horizon-Uno optimization profile") + add_lora_and_base_profiles( + builder, builder_config, ("token_id", "position_id"), cfg.max_block_size + ) + lora_rows = _lora_rows(network, token_id, lora_mask) attention_masks = add_active_prefix_causal_masks( network, @@ -465,6 +541,7 @@ def projection(lhs, logical: str, lhs_width: int, rhs_width: int, graph_name: st weights[f"{prefix}.{logical}_lora_a"], weights[f"{prefix}.{logical}_lora_b"], lora_mask, + lora_rows, lhs_width=lhs_width, rhs_width=rhs_width, rank=cfg.lora_rank, diff --git a/families/k2_horizon_uno/runtime/algorithm.cpp b/families/k2_horizon_uno/runtime/algorithm.cpp index f6d94f0895..c84a49dad0 100644 --- a/families/k2_horizon_uno/runtime/algorithm.cpp +++ b/families/k2_horizon_uno/runtime/algorithm.cpp @@ -313,6 +313,12 @@ std::vector k2_horizon_uno_draft_lora_mask(int32_t block_length) { return mask; } +bool k2_horizon_uno_mask_uses_lora(const std::vector& lora_mask) { + // NaN compares unequal to 0 and keeps the LoRA profile, as before. + return std::any_of(lora_mask.begin(), lora_mask.end(), + [](float value) { return value != 0.0F; }); +} + K2HorizonUnoCommitDecision k2_horizon_uno_decide_greedy_commit(const std::vector& proposal_tokens, const std::vector& verifier_next_tokens) { diff --git a/families/k2_horizon_uno/runtime/algorithm.h b/families/k2_horizon_uno/runtime/algorithm.h index d10cead01c..85d9b37b7a 100644 --- a/families/k2_horizon_uno/runtime/algorithm.h +++ b/families/k2_horizon_uno/runtime/algorithm.h @@ -54,6 +54,11 @@ std::vector k2_horizon_uno_deterministic_uniform_noise( // rows: [0, 1, ..., 1]. std::vector k2_horizon_uno_draft_lora_mask(int32_t block_length); +// True when any row has a nonzero LoRA mask. Calls whose mask is all 0 (prompt, +// verify and AR rows) run on the base profile, which skips the LoRA branch; +// its delta is multiplied by 0 there anyway. +bool k2_horizon_uno_mask_uses_lora(const std::vector& lora_mask); + // Greedy Psi-Spec verification. proposal_tokens is [clean, spec_1, ...]. // verifier_next_tokens contains the base-AR next token predicted after each // corresponding proposal row. A rejection commits the verified correction; diff --git a/families/k2_horizon_uno/runtime/pipeline.cpp b/families/k2_horizon_uno/runtime/pipeline.cpp index 56acbc1519..c846e11f52 100644 --- a/families/k2_horizon_uno/runtime/pipeline.cpp +++ b/families/k2_horizon_uno/runtime/pipeline.cpp @@ -27,10 +27,14 @@ int32_t checked_size_to_int32(std::size_t value, const char* label) { return static_cast(value); } -void validate_pipeline_components(const std::unique_ptr& decoder, +void validate_pipeline_components(const std::unique_ptr& lora_decoder, + const std::unique_ptr& base_decoder, const std::unique_ptr& cache) { - if (!decoder || !decoder->ok()) + if (!lora_decoder || !lora_decoder->ok() || !base_decoder || !base_decoder->ok()) throw std::runtime_error("K2HorizonUnoTextGenerationPipeline: invalid decoder module"); + if (lora_decoder->stream() != base_decoder->stream()) + throw std::runtime_error( + "K2HorizonUnoTextGenerationPipeline: decoders must share a stream"); if (!cache || !cache->ok()) throw std::runtime_error("K2HorizonUnoTextGenerationPipeline: invalid native KV cache"); } @@ -47,10 +51,20 @@ void validate_generation_inputs(const std::vector& token_ids, int32_t m } // namespace K2HorizonUnoTextGenerationPipeline::K2HorizonUnoTextGenerationPipeline( - std::unique_ptr decoder, std::unique_ptr cache, - std::shared_ptr tokenizer) - : decoder_(std::move(decoder)), cache_(std::move(cache)), tokenizer_(std::move(tokenizer)) { - validate_pipeline_components(decoder_, cache_); + std::unique_ptr lora_decoder, std::unique_ptr base_decoder, + std::unique_ptr cache, std::shared_ptr tokenizer) + : 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_); + // 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 + // count changes. Draft calls keep 8 rows (fewer only in the last cycle), + // so the LoRA context replays its graph. The base context changes row + // count in every request (prompt chunks of 8, k and 1 rows, then verify + // blocks); a recapture (2 to 13 ms) costs more than the replay saves on it + // (about 0.2 ms per call), so it runs without a graph. + lora_decoder_->enable_cuda_graph(); } TextResult K2HorizonUnoTextGenerationPipeline::generate(const std::string& prompt, @@ -212,8 +226,10 @@ K2HorizonUnoTextGenerationPipeline::generate_linear_psi_spec( void K2HorizonUnoTextGenerationPipeline::reset_generation_context() { const auto start = std::chrono::steady_clock::now(); cache_->reset(); - decoder_->reset_execution_context(); - cache_->bind_to(*decoder_); + lora_decoder_->reset_execution_context(); + base_decoder_->reset_execution_context(); + cache_->bind_to(*lora_decoder_); + cache_->bind_to(*base_decoder_); last_setup_ms_ = std::chrono::duration(std::chrono::steady_clock::now() - start).count(); } @@ -241,12 +257,17 @@ void K2HorizonUnoTextGenerationPipeline::run_block(const std::vector& t throw std::invalid_argument("K2-Horizon-Uno block input contract is invalid"); } const int32_t rows = checked_size_to_int32(token_ids.size(), "block length"); + // An all-zero mask adds nothing, so those calls skip the LoRA branch: the + // base profile takes an empty mask. + const bool uses_lora = k2_horizon_uno_mask_uses_lora(lora_mask); + ITrtModule& decoder = uses_lora ? *lora_decoder_ : *base_decoder_; TensorMap inputs; inputs["token_id"] = Tensor{const_cast(token_ids.data()), {rows}, DType::kInt32}; - inputs["lora_mask"] = Tensor{const_cast(lora_mask.data()), {rows}, DType::kFloat32}; + inputs["lora_mask"] = Tensor{ + const_cast(lora_mask.data()), {uses_lora ? rows : int64_t{0}}, DType::kFloat32}; cache_->prepare_block(inputs, rows); - const TensorMap outputs = decoder_->forward(inputs); + const TensorMap outputs = decoder.forward(inputs); const auto found = outputs.find("logits"); if (found == outputs.end() || found->second.dtype != DType::kFloat32 || found->second.shape != std::vector{rows, kVocabSize}) { diff --git a/families/k2_horizon_uno/runtime/pipeline.h b/families/k2_horizon_uno/runtime/pipeline.h index f79077b6b9..3452367abf 100644 --- a/families/k2_horizon_uno/runtime/pipeline.h +++ b/families/k2_horizon_uno/runtime/pipeline.h @@ -20,7 +20,11 @@ namespace trtmc { class K2HorizonUnoTextGenerationPipeline final : public ITextGeneration { public: - K2HorizonUnoTextGenerationPipeline(std::unique_ptr decoder, + // lora_decoder runs engine profile 0 (one LoRA mask value per row); + // base_decoder runs profile 1 (empty mask, LoRA branch skipped). Both are + // contexts of one engine and bind the same KV cache. + K2HorizonUnoTextGenerationPipeline(std::unique_ptr lora_decoder, + std::unique_ptr base_decoder, std::unique_ptr cache, std::shared_ptr tokenizer = nullptr); @@ -59,7 +63,8 @@ class K2HorizonUnoTextGenerationPipeline final : public ITextGeneration { std::vector effective_eos_ids(const TextGenerationConfig& cfg) const; void log_decode_receipt(const DecodeStats& stats) const; - std::unique_ptr decoder_; + std::unique_ptr lora_decoder_; + std::unique_ptr base_decoder_; 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 e6a4478a04..d8c308d325 100644 --- a/families/k2_horizon_uno/runtime/plugin.cpp +++ b/families/k2_horizon_uno/runtime/plugin.cpp @@ -47,9 +47,13 @@ std::string require_text_section(const BundleReader& bundle, std::string_view na return {data.begin(), data.end()}; } -std::unique_ptr load_engine(IBackend& backend, const std::vector& plan) { +constexpr std::int32_t kLoraProfile = 0; +constexpr std::int32_t kBaseProfile = 1; + +// One deserialized engine, one execution context per profile. +BackendDualProfileModules load_engine(IBackend& backend, const std::vector& plan) { const auto start = std::chrono::steady_clock::now(); - auto module = backend.create_module(plan.data(), plan.size(), {}); + auto modules = backend.create_dual_profile_modules(plan.data(), plan.size(), {}); const auto elapsed = std::chrono::duration(std::chrono::steady_clock::now() - start).count(); std::ostringstream receipt; @@ -57,10 +61,15 @@ std::unique_ptr load_engine(IBackend& backend, const std::vectorset_timing_label("k2_horizon_uno.engine.plan"); - return module; + if (modules.prefill->profile_idx() != kLoraProfile || + modules.decode->profile_idx() != kBaseProfile) + throw std::runtime_error("K2-Horizon-Uno engine contexts use unexpected profiles"); + modules.prefill->set_timing_label("k2_horizon_uno.engine.plan.lora"); + modules.decode->set_timing_label("k2_horizon_uno.engine.plan.base"); + return modules; } std::shared_ptr create_tokenizer(const BundleReader& bundle) { @@ -107,23 +116,34 @@ K2HorizonUnoKvCacheNames make_cache_names() { return names; } -void require_dynamic_vector_input(const ITrtModule& module, const std::string& name, DType dtype) { +void require_dynamic_vector_input(const ITrtModule& module, const std::string& name, DType dtype, + std::int32_t profile) { if (!module.has_input(name) || module.tensor_dtype(name) != dtype || !module.input_is_dynamic(name) || - module.input_profile_shape(name, 0, ProfileShapeSelector::kMin) != + module.input_profile_shape(name, profile, ProfileShapeSelector::kMin) != std::vector{1} || - module.input_profile_shape(name, 0, ProfileShapeSelector::kMax) != + module.input_profile_shape(name, profile, ProfileShapeSelector::kMax) != std::vector{kMaximumBlockLength}) { throw std::runtime_error("K2-Horizon-Uno dynamic input contract mismatch for '" + name + "'"); } - const auto optimum = module.input_profile_shape(name, 0, ProfileShapeSelector::kOpt); + 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 + "'"); } } +void require_empty_mask_profile(const ITrtModule& module, std::int32_t profile) { + for (const auto selector : + {ProfileShapeSelector::kMin, ProfileShapeSelector::kOpt, ProfileShapeSelector::kMax}) { + if (module.input_profile_shape("lora_mask", profile, selector) != + std::vector{0}) { + throw std::runtime_error("K2-Horizon-Uno base profile must take an empty lora_mask"); + } + } +} + std::set tensor_names(const std::vector& tensors) { std::set result; for (const auto& tensor : tensors) @@ -132,11 +152,15 @@ std::set tensor_names(const std::vector& tensors) { } void validate_engine(const ITrtModule& module, const K2HorizonUnoKvCacheNames& cache_names) { - if (module.optimization_profile_count() != 1) - throw std::runtime_error("K2-Horizon-Uno requires exactly one optimization profile"); - require_dynamic_vector_input(module, "token_id", DType::kInt32); - require_dynamic_vector_input(module, "position_id", DType::kInt32); - require_dynamic_vector_input(module, "lora_mask", DType::kFloat32); + if (module.optimization_profile_count() != 2) + throw std::runtime_error( + "K2-Horizon-Uno requires a LoRA profile and a base profile; rebuild the bundle"); + for (const auto profile : {kLoraProfile, kBaseProfile}) { + require_dynamic_vector_input(module, "token_id", DType::kInt32, profile); + require_dynamic_vector_input(module, "position_id", DType::kInt32, profile); + } + require_dynamic_vector_input(module, "lora_mask", DType::kFloat32, kLoraProfile); + require_empty_mask_profile(module, kBaseProfile); if (!module.has_output("logits") || module.tensor_dtype("logits") != DType::kFloat32 || module.tensor_shape("logits") != std::vector{kMaximumBlockLength, kVocabSize}) { @@ -203,17 +227,20 @@ ITask* create(const FamilyContext& context) { auto tokenizer = create_tokenizer(context.reader); validate_chat_tokenizer(*tokenizer); - auto decoder = load_engine(context.backend, require_section(context.reader, "engine.plan")); + auto decoders = load_engine(context.backend, require_section(context.reader, "engine.plan")); + auto lora_decoder = std::move(decoders.prefill); + auto base_decoder = std::move(decoders.decode); auto cache_names = make_cache_names(); - validate_engine(*decoder, cache_names); - auto cache = std::make_unique(max_cache_length, decoder->stream(), + validate_engine(*lora_decoder, cache_names); + auto cache = std::make_unique(max_cache_length, lora_decoder->stream(), std::move(cache_names)); if (!cache->ok()) throw std::runtime_error("K2-Horizon-Uno KV allocation failed"); - cache->bind_to(*decoder); + cache->bind_to(*lora_decoder); + cache->bind_to(*base_decoder); - return new K2HorizonUnoTextGenerationPipeline(std::move(decoder), std::move(cache), - std::move(tokenizer)); + return new K2HorizonUnoTextGenerationPipeline(std::move(lora_decoder), std::move(base_decoder), + std::move(cache), std::move(tokenizer)); } } // namespace trtmc::k2_horizon_uno 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 5abcb8c053..7635d7a1e9 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 @@ -142,6 +142,12 @@ void test_row_gate_is_exact() { for (std::size_t index = 1; index < mask.size(); ++index) noise_enabled = noise_enabled && mask[index] == 1.0F; check(noise_enabled, "all noise rows have LoRA enabled"); + check(trtmc::k2_horizon_uno_mask_uses_lora(mask), "draft mask selects the LoRA profile"); + check(trtmc::k2_horizon_uno_mask_uses_lora({0.0F, 0.0F, 0.5F}), + "any nonzero row selects the LoRA profile"); + check(!trtmc::k2_horizon_uno_mask_uses_lora(std::vector(8, 0.0F)), + "all-zero verify mask selects the base profile"); + check(!trtmc::k2_horizon_uno_mask_uses_lora({-0.0F}), "negative zero selects the base profile"); } void test_official_deterministic_uniform_golden() { diff --git a/families/k2_horizon_uno/tests/test_lora_profiles.py b/families/k2_horizon_uno/tests/test_lora_profiles.py new file mode 100644 index 0000000000..14c4c2ec60 --- /dev/null +++ b/families/k2_horizon_uno/tests/test_lora_profiles.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The LoRA profile applies the conditional branch; the base profile skips it.""" + +from __future__ import annotations + +import numpy as np +import pytest + +trt = pytest.importorskip("tensorrt") +torch = pytest.importorskip("torch") + +from families.k2_horizon_uno import model # noqa: E402 + + +_VOCAB, _IN, _OUT, _RANK, _SCALE, _ROWS = 16, 32, 48, 4, 2.0, 8 + + +def _bf16_bits(values: np.ndarray) -> np.ndarray: + tensor = torch.from_numpy(values.astype(np.float32)).to(torch.bfloat16) + return np.ascontiguousarray(tensor.view(torch.int16).numpy().view(np.uint16)) + + +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): + 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) + keepalive: list[np.ndarray] = [] + embedding = model._work_constant( + network, table.shape, table, work_dtype=trt.bfloat16, constant_keepalive=keepalive + ) + rows = network.add_gather(embedding, token_id, 0).get_output(0) + output = model._conditional_projection( + network, + rows, + weight, + lora_a, + lora_b, + lora_mask, + model._lora_rows(network, token_id, lora_mask), + lhs_width=_IN, + rhs_width=_OUT, + rank=_RANK, + scale=_SCALE, + work_dtype=trt.bfloat16, + constant_keepalive=keepalive, + name="probe", + ) + output = model._cast(network, output, trt.float32) + output.name = "output" + network.mark_output(output) + plan = builder.build_serialized_network(network, config) + assert plan is not None + engine = trt.Runtime(logger).deserialize_cuda_engine(plan) + assert engine.num_optimization_profiles == 2 + return engine + + +def _run(engine, profile: int, tokens: list[int], mask: list[float]) -> torch.Tensor: + stream = torch.cuda.Stream() + context = engine.create_execution_context() + context.set_optimization_profile_async(profile, stream.cuda_stream) + token_tensor = torch.tensor(tokens, dtype=torch.int32, device="cuda") + # An empty mask still needs a valid device address. + mask_tensor = torch.tensor([*mask, 0.0], dtype=torch.float32, device="cuda") + output = torch.full((len(tokens), _OUT), float("nan"), device="cuda") + context.set_input_shape("token_id", (len(tokens),)) + context.set_input_shape("lora_mask", (len(mask),)) + context.set_tensor_address("token_id", token_tensor.data_ptr()) + context.set_tensor_address("lora_mask", mask_tensor.data_ptr()) + context.set_tensor_address("output", output.data_ptr()) + stream.wait_stream(torch.cuda.current_stream()) + assert context.execute_async_v3(stream.cuda_stream) + stream.synchronize() + return output.cpu() + + +@pytest.mark.gpu +@pytest.mark.trt +def test_base_profile_skips_lora_and_lora_profile_applies_the_mask(): + rng = np.random.default_rng(0) + 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) + + tokens = [3, 1, 4, 1, 5, 9, 2, 6] + x = _from_bits(table)[tokens] + base = x @ _from_bits(weight) + delta = (x @ _from_bits(lora_a)) @ _from_bits(lora_b) * _SCALE + + draft_mask = [0.0] + [1.0] * 7 + lora = _run(engine, model.LORA_PROFILE, tokens, draft_mask) + expected = base + delta * torch.tensor(draft_mask)[:, None] + torch.testing.assert_close(lora, expected, rtol=2e-2, atol=2e-2) + assert not torch.allclose(lora[1:], base[1:], rtol=2e-2, atol=2e-2) + + zero_mask = _run(engine, model.LORA_PROFILE, tokens, [0.0] * 8) + for rows in (8, 3, 1): + skipped = _run(engine, model.BASE_PROFILE, tokens[:rows], []) + torch.testing.assert_close(skipped, base[:rows], rtol=2e-2, atol=2e-2) + if rows == 8: + assert torch.equal(skipped, zero_mask)