perf(k2_horizon_uno): skip the LoRA branch for all-0 masks and graph draft calls

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

families/k2_horizon_uno/model.py

@@ -120,13 +120,82 @@ 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,
     base_weight: np.ndarray,
     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,

families/k2_horizon_uno/runtime/algorithm.cpp

@@ -313,6 +313,12 @@ std::vector<float> k2_horizon_uno_draft_lora_mask(int32_t block_length) {
     return mask;
 }
 
+bool k2_horizon_uno_mask_uses_lora(const std::vector<float>& 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<int32_t>& proposal_tokens,
                                     const std::vector<int32_t>& verifier_next_tokens) {

families/k2_horizon_uno/runtime/algorithm.h

@@ -54,6 +54,11 @@ std::vector<int32_t> k2_horizon_uno_deterministic_uniform_noise(
 // rows: [0, 1, ..., 1].
 std::vector<float> 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<float>& 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;

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<int32_t>(value);
 }
 
-void validate_pipeline_components(const std::unique_ptr<ITrtModule>& decoder,
+void validate_pipeline_components(const std::unique_ptr<ITrtModule>& lora_decoder,
+                                  const std::unique_ptr<ITrtModule>& base_decoder,
                                   const std::unique_ptr<K2HorizonUnoKvCache>& 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<int32_t>& token_ids, int32_t m
 } // namespace
 
 K2HorizonUnoTextGenerationPipeline::K2HorizonUnoTextGenerationPipeline(
-    std::unique_ptr<ITrtModule> decoder, std::unique_ptr<K2HorizonUnoKvCache> cache,
-    std::shared_ptr<ITokenizer> tokenizer)
-    : decoder_(std::move(decoder)), cache_(std::move(cache)), tokenizer_(std::move(tokenizer)) {
-    validate_pipeline_components(decoder_, cache_);
+    std::unique_ptr<ITrtModule> lora_decoder, std::unique_ptr<ITrtModule> base_decoder,
+    std::unique_ptr<K2HorizonUnoKvCache> cache, std::shared_ptr<ITokenizer> 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<double, std::milli>(std::chrono::steady_clock::now() - start).count();
 }
@@ -241,12 +257,17 @@ void K2HorizonUnoTextGenerationPipeline::run_block(const std::vector<int32_t>& 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<int32_t*>(token_ids.data()), {rows}, DType::kInt32};
-    inputs["lora_mask"] = Tensor{const_cast<float*>(lora_mask.data()), {rows}, DType::kFloat32};
+    inputs["lora_mask"] = Tensor{
+        const_cast<float*>(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<int64_t>{rows, kVocabSize}) {

families/k2_horizon_uno/runtime/pipeline.h

@@ -20,7 +20,11 @@ namespace trtmc {
 
 class K2HorizonUnoTextGenerationPipeline final : public ITextGeneration {
   public:
-    K2HorizonUnoTextGenerationPipeline(std::unique_ptr<ITrtModule> 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<ITrtModule> lora_decoder,
+                                       std::unique_ptr<ITrtModule> base_decoder,
                                        std::unique_ptr<K2HorizonUnoKvCache> cache,
                                        std::shared_ptr<ITokenizer> tokenizer = nullptr);
 
@@ -59,7 +63,8 @@ class K2HorizonUnoTextGenerationPipeline final : public ITextGeneration {
     std::vector<int32_t> effective_eos_ids(const TextGenerationConfig& cfg) const;
     void log_decode_receipt(const DecodeStats& stats) const;
 
-    std::unique_ptr<ITrtModule> decoder_;
+    std::unique_ptr<ITrtModule> lora_decoder_;
+    std::unique_ptr<ITrtModule> base_decoder_;
     std::unique_ptr<K2HorizonUnoKvCache> cache_;
     std::shared_ptr<ITokenizer> tokenizer_;
     double last_setup_ms_{0.0};