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)