perf(phi4_multimodal): prefill and decode over the cached prefix, not the whole KV allocation

families/phi4_multimodal/default_decoder.py

@@ -35,6 +35,10 @@
     from .checkpoint_mapper import WeightDict
     from typing import Any as QuantContext
 
+# Optimal cache rows of a runtime-sized decode engine: a typical image prompt
+# plus its answer.
+DECODE_OPT_CACHE_ROWS = 2048
+
 
 def _mark_debug_output(
     network: trt.INetworkDefinition,
@@ -164,6 +168,15 @@ def build_standard_decoder_engine(
     kv_attention_size = graph_blocks.infer_kv_attention_size(
         weights, num_kv_heads=num_kv_heads, head_dim=head_dim)
     attention_window = max_cache_length + 1
+    # The split layout's decode engine attends over the cache rows the runtime
+    # binds (the cached prefix, 1..max_cache_length) instead of the whole
+    # allocation; its inputs are cache_k_i/cache_v_i [rows, kv] and an
+    # attention_mask [1, rows + 1].
+    dynamic_cache_rows = decoder_engine_role == "decode"
+    if dynamic_cache_rows and position_type == "alibi":
+        raise NotImplementedError(
+            "a runtime-sized decode cache needs RoPE or learned positions")
+    cache_rows = -1 if dynamic_cache_rows else max_cache_length
 
     builder_context = create_builder_context(
         verbose=verbose,
@@ -191,7 +204,7 @@ def build_standard_decoder_engine(
     position_id = network.add_input("position_id", trt.int32, (1,))
     attention_mask = network.add_input(
         "attention_mask", trt.float32,
-        (1, attention_window))
+        (1, -1 if dynamic_cache_rows else attention_window))
 
     # Optional VL inputs: when embed_input=True, the decoder can accept
     # a pre-computed embedding vector instead of a token ID.
@@ -209,14 +222,25 @@ def build_standard_decoder_engine(
         ck = network.add_input(
             graph_ops.layer_tensor_name("cache_k", i),
             work_trt_dtype,
-            (max_cache_length, kv_attention_size))
+            (cache_rows, kv_attention_size))
         cv = network.add_input(
             graph_ops.layer_tensor_name("cache_v", i),
             work_trt_dtype,
-            (max_cache_length, kv_attention_size))
+            (cache_rows, kv_attention_size))
         cache_k_inputs.append(ck)
         cache_v_inputs.append(cv)
 
+    if dynamic_cache_rows:
+        opt_rows = min(max_cache_length, DECODE_OPT_CACHE_ROWS)
+        profile = builder.create_optimization_profile()
+        profile.set_shape(
+            "attention_mask", (1, 2), (1, opt_rows + 1), (1, attention_window))
+        for cache_input in (*cache_k_inputs, *cache_v_inputs):
+            profile.set_shape(
+                cache_input.name, (1, kv_attention_size),
+                (opt_rows, kv_attention_size), (max_cache_length, kv_attention_size))
+        trt_config.add_optimization_profile(profile)
+
 
     # Cast the attention mask so elementwise operands share the work dtype
     if work_trt_dtype != trt.float32:
@@ -401,6 +425,7 @@ def _cast_layer_dtype(tensor: trt.ITensor | None) -> trt.ITensor | None:
             sin_half_tensor=_cast_layer_dtype(sin_half_tensor),
             rotary_embedding_dim=rotary_embedding_dim,
             interleaved_rope=interleaved_rope,
+            dynamic_cache_rows=dynamic_cache_rows,
         )
 
         hidden_state = _cast_work_dtype(result["hidden"])
@@ -531,6 +556,7 @@ def _add_decoder_layer(
     rotary_embedding_dim: int = 0,
     interleaved_rope: bool = False,
     eps: float | None = None,
+    dynamic_cache_rows: bool = False,
 ) -> dict[str, trt.ITensor]:
     """Add one standard decoder layer block. Returns hidden, present_k, present_v."""
 
@@ -554,6 +580,7 @@ def _add_decoder_layer(
         sin_half_tensor=sin_half_tensor,
         rotary_embedding_dim=rotary_embedding_dim,
         interleaved_rope=interleaved_rope,
+        dynamic_cache_rows=dynamic_cache_rows,
     )
     attn_out = attn["attn_out"]
     present_k = attn["present_k"]

families/phi4_multimodal/default_dual_profile_decoder.py

@@ -30,9 +30,9 @@
   Inputs (dynamic shapes — Sq varies by profile)
     token_id        int32   (-1,)
     position_id     int32   (-1,)
-    attention_mask  float32 (-1, -1)                 # (Sq, max_cache + Sq)
-    cache_k_i       fp16/f32 (max_cache, kv_size)    # static
-    cache_v_i       fp16/f32 (max_cache, kv_size)    # static
+    attention_mask  float32 (-1, -1)                 # (Sq, cache_rows + Sq)
+    cache_k_i       fp16/f32 (cache_rows, kv_size)   # max_cache; 1..max_cache (prefill engine)
+    cache_v_i       fp16/f32 (cache_rows, kv_size)   # max_cache; 1..max_cache (prefill engine)
   Outputs
     logits          float32 (1, vocab)               # last-row sliced inside the engine
     present_k_i     fp16/f32 (-1, kv_size)           # (Sq, kv_size)
@@ -231,8 +231,13 @@ def build_dual_profile_decoder_engine(
         use_input_embed_tensor = network.add_input(
             "use_input_embed", trt.float32, (-1, 1))
 
+    # A prefill engine takes runtime-sized cache inputs: the runtime binds only
+    # the cached prefix (one masked row for an empty cache), so attention
+    # follows the context rather than the allocated capacity. Dual-profile
+    # engines keep the static cache that their decode profile shares.
+    dynamic_cache_rows = profile_mode == "prefill"
     cache_shape: tuple[int, int]
-    cache_shape = (max_cache_length, kv_attention_size)
+    cache_shape = (-1 if dynamic_cache_rows else max_cache_length, kv_attention_size)
     cache_k_inputs: list[trt.ITensor] = []
     cache_v_inputs: list[trt.ITensor] = []
     for i in range(num_layers):
@@ -258,11 +263,21 @@ def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False):
         min_sq = opt_sq if fixed else 1
         prof.set_shape("token_id", (min_sq,), (opt_sq,), (max_sq,))
         prof.set_shape("position_id", (min_sq,), (opt_sq,), (max_sq,))
+        # Prompts are prefilled from an empty cache, so its optimum is one row.
+        min_rows, opt_rows = (1, 1) if dynamic_cache_rows else (max_cache_length,) * 2
         prof.set_shape(
             "attention_mask",
-            (min_sq, max_cache_length + min_sq),
-            (opt_sq, max_cache_length + opt_sq),
+            (min_sq, min_rows + min_sq),
+            (opt_sq, opt_rows + opt_sq),
             (max_sq, max_cache_length + max_sq))
+        if dynamic_cache_rows:
+            for layer in range(num_layers):
+                for name in ("cache_k", "cache_v"):
+                    prof.set_shape(
+                        graph_ops.layer_tensor_name(name, layer),
+                        (min_rows, kv_attention_size),
+                        (opt_rows, kv_attention_size),
+                        (max_cache_length, kv_attention_size))
         if embed_input:
             prof.set_shape(
                 "input_embed", (min_sq, hidden), (opt_sq, hidden), (max_sq, hidden))

families/phi4_multimodal/graph_blocks.py

@@ -146,9 +146,13 @@ def add_attention_block(
     rotary_embedding_dim: int = 0,
     interleaved_rope: bool = False,
     sequence_length: int | None = 1,
+    dynamic_cache_rows: bool = False,
 ) -> dict[str, trt.ITensor]:
     """Pre-norm -> QKV -> RoPE -> cache concat -> attention -> output proj.
 
+    ``dynamic_cache_rows``: the cache inputs have a runtime row count, so the
+    attention reads ``rows + sequence_length`` keys instead of the whole window.
+
     Returns {"normed": ..., "attn_out": ..., "present_k": ..., "present_v": ...}.
     Does NOT apply residual -- callers compose the residual pattern.
 
@@ -253,7 +257,7 @@ def add_attention_block(
     # ------------------------------------------------------------------ #
     # Attention core — native IAttention                                  #
     # ------------------------------------------------------------------ #
-    kv_seq = None if sequence_length is None else attention_window
+    kv_seq = None if sequence_length is None or dynamic_cache_rows else attention_window
     if alibi_slopes_tensor is not None:
         if alibi_indices_tensor is None:
             raise ValueError("ALiBi attention requires cache position indices")

families/phi4_multimodal/runtime/kv_cache.cpp

@@ -71,20 +71,22 @@ void Phi4MultimodalKvCache::write_position_input(TensorMap& inputs, int32_t seq_
 }
 
 void Phi4MultimodalKvCache::write_batched_mask(TensorMap& inputs, int32_t seq_len) {
-    // Batched prefill mask: (seq_len, max_length + seq_len). Columns
-    // [0, valid) are visible cache, [valid, max_length) are stale slots,
-    // [max_length, max_length+seq_len) are the new tokens — causal so
-    // token i sees tokens 0..i.
-    const int32_t valid = std::max(0, std::min(position_, max_length_));
-    const int32_t kv_len = max_length_ + seq_len;
+    // Batched prefill mask: (seq_len, cache_rows + seq_len), cache_rows the
+    // bound cache prefix (max_length for a static engine). Columns [0, valid)
+    // are visible cache, [valid, cache_rows) are stale slots, and
+    // [cache_rows, cache_rows+seq_len) are the new tokens — causal so token i
+    // sees tokens 0..i.
+    const int32_t cache_rows = prefill_cache_rows();
+    const int32_t valid = std::max(0, std::min(position_, cache_rows));
+    const int32_t kv_len = cache_rows + seq_len;
     const std::size_t total = static_cast<std::size_t>(seq_len) * static_cast<std::size_t>(kv_len);
     mask_buf_.assign(total, kMaskedScore);
     for (int32_t i = 0; i < seq_len; ++i) {
         const std::size_t row = static_cast<std::size_t>(i) * static_cast<std::size_t>(kv_len);
         for (int32_t j = 0; j < valid; ++j)
             mask_buf_[row + static_cast<std::size_t>(j)] = 0.0f;
         for (int32_t j = 0; j <= i; ++j)
-            mask_buf_[row + static_cast<std::size_t>(max_length_) + static_cast<std::size_t>(j)] =
+            mask_buf_[row + static_cast<std::size_t>(cache_rows) + static_cast<std::size_t>(j)] =
                 0.0f;
     }
     Tensor mask_t;
@@ -98,16 +100,17 @@ void Phi4MultimodalKvCache::write_bidirectional_mask(TensorMap& inputs, int32_t
     // Diffusion block mask: all valid prefix cache rows are visible, stale cache
     // rows are hidden, and every token in the current block can see every other
     // token in the current block.
-    const int32_t valid = std::max(0, std::min(position_, max_length_));
-    const int32_t kv_len = max_length_ + seq_len;
+    const int32_t cache_rows = prefill_cache_rows();
+    const int32_t valid = std::max(0, std::min(position_, cache_rows));
+    const int32_t kv_len = cache_rows + seq_len;
     const std::size_t total = static_cast<std::size_t>(seq_len) * static_cast<std::size_t>(kv_len);
     mask_buf_.assign(total, kMaskedScore);
     for (int32_t i = 0; i < seq_len; ++i) {
         const std::size_t row = static_cast<std::size_t>(i) * static_cast<std::size_t>(kv_len);
         for (int32_t j = 0; j < valid; ++j)
             mask_buf_[row + static_cast<std::size_t>(j)] = 0.0f;
         for (int32_t j = 0; j < seq_len; ++j)
-            mask_buf_[row + static_cast<std::size_t>(max_length_) + static_cast<std::size_t>(j)] =
+            mask_buf_[row + static_cast<std::size_t>(cache_rows) + static_cast<std::size_t>(j)] =
                 0.0f;
     }
     Tensor mask_t;
@@ -124,8 +127,16 @@ void Phi4MultimodalKvCache::write_decode_mask(TensorMap& inputs) {
     if (mask_rank != 2)
         throw std::runtime_error("phi4_multimodal decoder has an invalid attention-mask rank");
 
+    // Columns [0, valid) are the cached prefix, [valid, cache_rows) stale rows
+    // and the last column the new token; a runtime-sized decode engine binds
+    // only the bucketed prefix instead of all max_length rows.
     const int32_t valid = std::max(0, std::min(position_, max_length_));
-    const int32_t mask_width = max_length_ + 1;
+    int32_t cache_rows = max_length_;
+    if (decode_dynamic_rows_) {
+        cache_rows = bucketed_cache_rows();
+        bind_decode_cache_rows(cache_rows);
+    }
+    const int32_t mask_width = cache_rows + 1;
     if (mask_buf_.size() != static_cast<std::size_t>(mask_width))
         mask_buf_.resize(static_cast<std::size_t>(mask_width));
     std::fill(mask_buf_.begin(), mask_buf_.end(), kMaskedScore);
@@ -162,10 +173,12 @@ void Phi4MultimodalKvCache::bind_to(ITrtModule& module) {
         throw std::runtime_error("phi4_multimodal KV cache has no cache tensor names");
     const auto cache_shape = module.tensor_shape(names_.cache_k.front());
     const std::vector<int64_t> expected_cache_shape{max_length_, kv_dim_};
-    if (module.input_rank(names_.cache_k.front()) != 2 ||
-        module.input_is_dynamic(names_.cache_k.front()) || cache_shape != expected_cache_shape) {
-        throw std::runtime_error(
-            "phi4_multimodal decoder does not match its fixed KV-cache contract");
+    const bool dynamic_rows = module.input_is_dynamic(names_.cache_k.front());
+    const bool shape_matches = dynamic_rows
+                                   ? cache_shape.size() == 2 && cache_shape.back() == kv_dim_
+                                   : cache_shape == expected_cache_shape;
+    if (module.input_rank(names_.cache_k.front()) != 2 || !shape_matches) {
+        throw std::runtime_error("phi4_multimodal decoder does not match its KV-cache contract");
     }
     const int32_t mask_rank = module.input_rank(names_.attention_mask);
     if (mask_rank != 2)
@@ -180,17 +193,52 @@ void Phi4MultimodalKvCache::bind_to(ITrtModule& module) {
         module.bind_external(names_.present_k[layer], present_k_[layer].data());
         module.bind_external(names_.present_v[layer], present_v_[layer].data());
     }
+    decode_dynamic_rows_ = dynamic_rows;
+    decode_cache_rows_ = 0;
+    if (decode_dynamic_rows_)
+        bind_decode_cache_rows(bucketed_cache_rows());
+}
+
+int32_t Phi4MultimodalKvCache::bucketed_cache_rows() const {
+    const int32_t used = std::max(0, std::min(position_, max_length_));
+    const int32_t bucketed = (used + kCacheRowBucket - 1) / kCacheRowBucket * kCacheRowBucket;
+    return std::max(1, std::min(bucketed, max_length_));
+}
+
+void Phi4MultimodalKvCache::bind_decode_cache_rows(int32_t rows) {
+    if (rows == decode_cache_rows_)
+        return;
+    const std::vector<int64_t> cache_shape{rows, kv_dim_};
+    for (int32_t i = 0; i < num_layers_; ++i) {
+        const auto layer = static_cast<std::size_t>(i);
+        bound_module_->bind_external(names_.cache_k[layer], cache_k_[layer].data(), cache_shape);
+        bound_module_->bind_external(names_.cache_v[layer], cache_v_[layer].data(), cache_shape);
+    }
+    decode_cache_rows_ = rows;
 }
 
 void Phi4MultimodalKvCache::bind_cache_inputs(ITrtModule& module) {
     has_position_input_ = module.has_input(names_.position_id);
+    prefill_cache_rows_ = 0;
+    if (!names_.cache_k.empty() && module.input_is_dynamic(names_.cache_k.front()))
+        prefill_cache_rows_ = bucketed_cache_rows();
+    const std::vector<int64_t> cache_shape{prefill_cache_rows_, kv_dim_};
     for (int32_t i = 0; i < num_layers_; ++i) {
         auto li = static_cast<std::size_t>(i);
-        module.bind_external(names_.cache_k[li], cache_k_[li].data());
-        module.bind_external(names_.cache_v[li], cache_v_[li].data());
+        if (prefill_cache_rows_ > 0) {
+            module.bind_external(names_.cache_k[li], cache_k_[li].data(), cache_shape);
+            module.bind_external(names_.cache_v[li], cache_v_[li].data(), cache_shape);
+        } else {
+            module.bind_external(names_.cache_k[li], cache_k_[li].data());
+            module.bind_external(names_.cache_v[li], cache_v_[li].data());
+        }
     }
 }
 
+int32_t Phi4MultimodalKvCache::prefill_cache_rows() const {
+    return prefill_cache_rows_ > 0 ? prefill_cache_rows_ : max_length_;
+}
+
 void Phi4MultimodalKvCache::write_prefill_kv(const std::vector<const void*>& prefill_k,
                                              const std::vector<const void*>& prefill_v,
                                              int32_t seq_len) {
@@ -290,6 +338,7 @@ void Phi4MultimodalKvCache::reset() {
     // Reset only the logical sequence length. Attention masks hide every
     // stale cache row, and each present row is overwritten before use.
     position_ = 0;
+    prefill_cache_rows_ = 0;
 }
 
 std::size_t Phi4MultimodalKvCache::device_memory_bytes() const {

families/phi4_multimodal/runtime/kv_cache.h

@@ -85,14 +85,29 @@ class Phi4MultimodalKvCache : public Phi4MultimodalInferenceState {
     // prefill ITrtModule whose present_k/v outputs have shape (Sq, kv_dim)
     // — too big for Phi4MultimodalKvCache's single-row present buffer. The caller reads
     // the prefill outputs directly from the module's own allocations and
-    // copies them via write_prefill_kv().
+    // copies them via write_prefill_kv(). A prefill engine with runtime-sized
+    // cache inputs gets only the cached prefix, rounded up to kCacheRowBucket
+    // rows (one masked row for an empty cache), so its attention follows the
+    // context rather than the allocated capacity. A runtime-sized decode
+    // engine (bind_to) is bound the same way before every step.
     void bind_cache_inputs(ITrtModule& module);
 
+    // Cache rows of the next batched (prefill) step: the rows bound by
+    // bind_cache_inputs() for a runtime-sized prefill engine, else max_length.
+    int32_t prefill_cache_rows() const;
+
+    static constexpr int32_t kCacheRowBucket = 256;
+
   private:
     void write_position_input(TensorMap& inputs, int32_t seq_len);
     void write_batched_mask(TensorMap& inputs, int32_t seq_len);
     void write_bidirectional_mask(TensorMap& inputs, int32_t seq_len);
     void write_decode_mask(TensorMap& inputs);
+    // The cached prefix rounded up to kCacheRowBucket rows, 1..max_length.
+    int32_t bucketed_cache_rows() const;
+    // Bind `rows` cache rows to the runtime-sized decode engine (no-op when
+    // they are already bound).
+    void bind_decode_cache_rows(int32_t rows);
 
     std::vector<DeviceTensor> cache_k_;   // [num_layers], shape [max_length, kv_dim]
     std::vector<DeviceTensor> cache_v_;   // [num_layers]
@@ -111,6 +126,12 @@ class Phi4MultimodalKvCache : public Phi4MultimodalInferenceState {
     std::size_t cache_element_size_{sizeof(float)};
     Phi4MultimodalKvCacheNames names_;
     ITrtModule* bound_module_{nullptr};
+    // Rows bound on a runtime-sized prefill engine; 0 when its cache is static.
+    int32_t prefill_cache_rows_{0};
+    // The bound decode engine takes runtime-sized cache inputs, and the rows
+    // currently bound to it (0 before the first step).
+    bool decode_dynamic_rows_{false};
+    int32_t decode_cache_rows_{0};
 };
 
 } // namespace trtmc