From 4d72959cfcae56da6dc940253ff514422b7b0734 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Tue, 10 Nov 2026 18:59:47 +0000 Subject: [PATCH 1/2] perf(phi4_multimodal): prefill over the cached prefix, not the whole KV allocation The prefill engine's cache inputs are runtime-sized (1..max_cache rows, optimum 1) and the runtime binds only the cached prefix rounded up to 256 rows, one masked row for an empty cache, with the batched mask sized to match. An image prompt, prefilled from an empty cache, no longer attends over all 10368 rows of the catalog bundle's allocation. The decode engine is unchanged. Signed-off-by: chaofengw --- .../default_dual_profile_decoder.py | 27 +++- families/phi4_multimodal/runtime/kv_cache.cpp | 45 ++++-- families/phi4_multimodal/runtime/kv_cache.h | 13 +- .../tests/test_prefill_cache_rows.py | 138 ++++++++++++++++++ 4 files changed, 204 insertions(+), 19 deletions(-) create mode 100644 families/phi4_multimodal/tests/test_prefill_cache_rows.py diff --git a/families/phi4_multimodal/default_dual_profile_decoder.py b/families/phi4_multimodal/default_dual_profile_decoder.py index 5a039da5e8..c725817d0e 100644 --- a/families/phi4_multimodal/default_dual_profile_decoder.py +++ b/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)) diff --git a/families/phi4_multimodal/runtime/kv_cache.cpp b/families/phi4_multimodal/runtime/kv_cache.cpp index 0c98c8d9a0..12ad25b1ac 100644 --- a/families/phi4_multimodal/runtime/kv_cache.cpp +++ b/families/phi4_multimodal/runtime/kv_cache.cpp @@ -71,12 +71,14 @@ 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(seq_len) * static_cast(kv_len); mask_buf_.assign(total, kMaskedScore); for (int32_t i = 0; i < seq_len; ++i) { @@ -84,7 +86,7 @@ void Phi4MultimodalKvCache::write_batched_mask(TensorMap& inputs, int32_t seq_le for (int32_t j = 0; j < valid; ++j) mask_buf_[row + static_cast(j)] = 0.0f; for (int32_t j = 0; j <= i; ++j) - mask_buf_[row + static_cast(max_length_) + static_cast(j)] = + mask_buf_[row + static_cast(cache_rows) + static_cast(j)] = 0.0f; } Tensor mask_t; @@ -98,8 +100,9 @@ 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(seq_len) * static_cast(kv_len); mask_buf_.assign(total, kMaskedScore); for (int32_t i = 0; i < seq_len; ++i) { @@ -107,7 +110,7 @@ void Phi4MultimodalKvCache::write_bidirectional_mask(TensorMap& inputs, int32_t for (int32_t j = 0; j < valid; ++j) mask_buf_[row + static_cast(j)] = 0.0f; for (int32_t j = 0; j < seq_len; ++j) - mask_buf_[row + static_cast(max_length_) + static_cast(j)] = + mask_buf_[row + static_cast(cache_rows) + static_cast(j)] = 0.0f; } Tensor mask_t; @@ -184,13 +187,30 @@ void Phi4MultimodalKvCache::bind_to(ITrtModule& module) { 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())) { + const int32_t used = std::max(0, std::min(position_, max_length_)); + const int32_t bucketed = + (used + kPrefillCacheRowBucket - 1) / kPrefillCacheRowBucket * kPrefillCacheRowBucket; + prefill_cache_rows_ = std::max(1, std::min(bucketed, max_length_)); + } + const std::vector cache_shape{prefill_cache_rows_, kv_dim_}; for (int32_t i = 0; i < num_layers_; ++i) { auto li = static_cast(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& prefill_k, const std::vector& prefill_v, int32_t seq_len) { @@ -290,6 +310,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 { diff --git a/families/phi4_multimodal/runtime/kv_cache.h b/families/phi4_multimodal/runtime/kv_cache.h index 54d85289b2..9069b8e552 100644 --- a/families/phi4_multimodal/runtime/kv_cache.h +++ b/families/phi4_multimodal/runtime/kv_cache.h @@ -85,9 +85,18 @@ 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 + // kPrefillCacheRowBucket rows (one masked row for an empty cache), so its + // attention follows the context rather than the allocated capacity. 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 kPrefillCacheRowBucket = 256; + private: void write_position_input(TensorMap& inputs, int32_t seq_len); void write_batched_mask(TensorMap& inputs, int32_t seq_len); @@ -111,6 +120,8 @@ 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}; }; } // namespace trtmc diff --git a/families/phi4_multimodal/tests/test_prefill_cache_rows.py b/families/phi4_multimodal/tests/test_prefill_cache_rows.py new file mode 100644 index 0000000000..238e7f639c --- /dev/null +++ b/families/phi4_multimodal/tests/test_prefill_cache_rows.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The prefill engine attends over the bound cache prefix, not the whole allocation.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +trt = pytest.importorskip("tensorrt") +torch = pytest.importorskip("torch") + +from families.phi4_multimodal import model as family_model # noqa: E402 +from families.phi4_multimodal.config import ModelConfig # noqa: E402 + +from .test_family_contract import _rand, _write_config, _write_safetensors # noqa: E402 + +MAX_CACHE = 64 +MASKED = -1.0e4 +VOCAB, HIDDEN, HEADS, KV_HEADS, MLP = 64, 32, 4, 2, 64 + + +def _write_checkpoint(path: Path) -> ModelConfig: + head_dim = HIDDEN // HEADS + _write_config( + path, + { + "model_type": "phi4mm", + "vocab_size": VOCAB, + "hidden_size": HIDDEN, + "num_hidden_layers": 2, + "num_attention_heads": HEADS, + "num_key_value_heads": KV_HEADS, + "intermediate_size": MLP, + "max_position_embeddings": 4 * MAX_CACHE, + "rope_theta": 10000.0, + "rms_norm_eps": 1e-5, + }, + ) + tensors = {"model.embed_tokens.weight": _rand(VOCAB, HIDDEN) * 0.1} + for layer in range(2): + prefix = f"model.layers.{layer}" + qkv = HEADS * head_dim + 2 * KV_HEADS * head_dim + tensors.update( + { + f"{prefix}.input_layernorm.weight": np.ones(HIDDEN, np.float32), + f"{prefix}.post_attention_layernorm.weight": np.ones(HIDDEN, np.float32), + f"{prefix}.self_attn.qkv_proj.base_layer.weight": _rand(qkv, HIDDEN) * 0.1, + f"{prefix}.self_attn.o_proj.base_layer.weight": _rand(HIDDEN, HIDDEN) * 0.1, + f"{prefix}.mlp.gate_up_proj.base_layer.weight": _rand(2 * MLP, HIDDEN) * 0.1, + f"{prefix}.mlp.down_proj.base_layer.weight": _rand(HIDDEN, MLP) * 0.1, + } + ) + tensors["model.norm.weight"] = np.ones(HIDDEN, np.float32) + tensors["lm_head.weight"] = _rand(VOCAB, HIDDEN) * 0.1 + _write_safetensors(path, tensors) + return ModelConfig.from_dir(path) + + +def _run(engine, cache_rows: int, caches: dict, tokens: list[int], position: int) -> dict: + """One prefill launch binding the first ``cache_rows`` rows of each cache.""" + context = engine.create_execution_context() + stream = torch.cuda.current_stream().cuda_stream + rows = len(tokens) + mask = torch.full((rows, cache_rows + rows), MASKED) + mask[:, :position] = 0.0 + for row in range(rows): + mask[row, cache_rows : cache_rows + row + 1] = 0.0 + inputs = { + "token_id": torch.tensor(tokens, dtype=torch.int32), + "position_id": torch.arange(position, position + rows, dtype=torch.int32), + "attention_mask": mask, + "input_embed": torch.zeros(rows, HIDDEN), + "use_input_embed": torch.zeros(rows, 1), + } + names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)] + bound = {} + for name in names: + if engine.get_tensor_mode(name) != trt.TensorIOMode.INPUT: + continue + if name in caches: + value = caches[name][:cache_rows] + else: + value = inputs[name].cuda() + bound[name] = value + if -1 in tuple(engine.get_tensor_shape(name)): + assert context.set_input_shape(name, tuple(value.shape)), name + context.set_tensor_address(name, value.data_ptr()) + outputs = {} + for name in names: + if engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT: + outputs[name] = torch.empty( + tuple(context.get_tensor_shape(name)), dtype=torch.float32, device="cuda" + ) + context.set_tensor_address(name, outputs[name].data_ptr()) + assert context.execute_async_v3(stream) + torch.cuda.synchronize() + return {name: value.cpu().numpy() for name, value in outputs.items()} + + +@pytest.mark.gpu +@pytest.mark.trt +def test_prefill_over_the_cached_prefix_matches_the_whole_cache(tmp_path: Path) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + config = _write_checkpoint(tmp_path) + config.raw["_model_dir"] = str(tmp_path) + config.raw["_decoder_engine_role"] = "prefill" + instance = family_model._Phi4MultimodalModel() + weights = instance.load_weights(str(tmp_path), config) + engine = trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine( + instance.build_engine(config, weights, MAX_CACHE, precision="fp32") + ) + assert tuple(engine.get_tensor_shape("cache_k_0"))[0] == -1 + assert [tuple(s) for s in engine.get_tensor_profile_shape("cache_k_0", 0)] == [ + (1, engine.get_tensor_shape("cache_k_0")[1]), + (1, engine.get_tensor_shape("cache_k_0")[1]), + (MAX_CACHE, engine.get_tensor_shape("cache_k_0")[1]), + ] + + caches = {} + for i in range(engine.num_io_tensors): + name = engine.get_tensor_name(i) + if name.startswith("cache_"): + width = engine.get_tensor_shape(name)[1] + # Rows beyond the prefix hold garbage that the mask must hide. + caches[name] = torch.randn(MAX_CACHE, width, device="cuda") + + rng = np.random.default_rng(0) + for position, rows in ((0, 1), (5, 16)): + tokens = rng.integers(0, VOCAB, 7).tolist() + whole = _run(engine, MAX_CACHE, caches, tokens, position) + prefix = _run(engine, rows, caches, tokens, position) + for name, value in whole.items(): + np.testing.assert_allclose(prefix[name], value, rtol=1e-4, atol=1e-4, err_msg=name) From d8d899fbf6c4842206546898fba4b73d92496eb4 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Tue, 10 Nov 2026 20:17:26 +0000 Subject: [PATCH 2/2] perf(phi4_multimodal): decode over the cached prefix, not the whole KV allocation The split layout's one-token decode engine took cache_k_i/cache_v_i as [max_cache, kv] and attended over all 10368 rows at every step. It now takes [rows, kv] (profile 1 / 2048 / max_cache) and an attention_mask [1, rows + 1], and the runtime binds the cached prefix rounded up to 256 rows before each step, rebinding only when the bucket changes. Bundles with a static decode engine keep the full-width binding. Signed-off-by: chaofengw --- families/phi4_multimodal/default_decoder.py | 33 +++++++++- families/phi4_multimodal/graph_blocks.py | 6 +- families/phi4_multimodal/runtime/kv_cache.cpp | 50 ++++++++++---- families/phi4_multimodal/runtime/kv_cache.h | 18 +++-- .../tests/test_prefill_cache_rows.py | 65 +++++++++++++++---- 5 files changed, 139 insertions(+), 33 deletions(-) diff --git a/families/phi4_multimodal/default_decoder.py b/families/phi4_multimodal/default_decoder.py index 2bdeaca7de..f6fccc3a81 100644 --- a/families/phi4_multimodal/default_decoder.py +++ b/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"] diff --git a/families/phi4_multimodal/graph_blocks.py b/families/phi4_multimodal/graph_blocks.py index c28a9bb12a..cd44db6feb 100644 --- a/families/phi4_multimodal/graph_blocks.py +++ b/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") diff --git a/families/phi4_multimodal/runtime/kv_cache.cpp b/families/phi4_multimodal/runtime/kv_cache.cpp index 12ad25b1ac..7a3cf916d1 100644 --- a/families/phi4_multimodal/runtime/kv_cache.cpp +++ b/families/phi4_multimodal/runtime/kv_cache.cpp @@ -127,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(mask_width)) mask_buf_.resize(static_cast(mask_width)); std::fill(mask_buf_.begin(), mask_buf_.end(), kMaskedScore); @@ -165,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 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) @@ -183,17 +193,35 @@ 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 cache_shape{rows, kv_dim_}; + for (int32_t i = 0; i < num_layers_; ++i) { + const auto layer = static_cast(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())) { - const int32_t used = std::max(0, std::min(position_, max_length_)); - const int32_t bucketed = - (used + kPrefillCacheRowBucket - 1) / kPrefillCacheRowBucket * kPrefillCacheRowBucket; - prefill_cache_rows_ = std::max(1, std::min(bucketed, max_length_)); - } + if (!names_.cache_k.empty() && module.input_is_dynamic(names_.cache_k.front())) + prefill_cache_rows_ = bucketed_cache_rows(); const std::vector cache_shape{prefill_cache_rows_, kv_dim_}; for (int32_t i = 0; i < num_layers_; ++i) { auto li = static_cast(i); diff --git a/families/phi4_multimodal/runtime/kv_cache.h b/families/phi4_multimodal/runtime/kv_cache.h index 9069b8e552..297b14a3b1 100644 --- a/families/phi4_multimodal/runtime/kv_cache.h +++ b/families/phi4_multimodal/runtime/kv_cache.h @@ -86,22 +86,28 @@ class Phi4MultimodalKvCache : public Phi4MultimodalInferenceState { // — 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(). A prefill engine with runtime-sized - // cache inputs gets only the cached prefix, rounded up to - // kPrefillCacheRowBucket rows (one masked row for an empty cache), so its - // attention follows the context rather than the allocated capacity. + // 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 kPrefillCacheRowBucket = 256; + 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 cache_k_; // [num_layers], shape [max_length, kv_dim] std::vector cache_v_; // [num_layers] @@ -122,6 +128,10 @@ class Phi4MultimodalKvCache : public Phi4MultimodalInferenceState { 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 diff --git a/families/phi4_multimodal/tests/test_prefill_cache_rows.py b/families/phi4_multimodal/tests/test_prefill_cache_rows.py index 238e7f639c..a786000627 100644 --- a/families/phi4_multimodal/tests/test_prefill_cache_rows.py +++ b/families/phi4_multimodal/tests/test_prefill_cache_rows.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The prefill engine attends over the bound cache prefix, not the whole allocation.""" +"""The prefill and decode engines attend over the bound cache prefix, not the whole allocation.""" from __future__ import annotations @@ -13,6 +13,7 @@ trt = pytest.importorskip("tensorrt") torch = pytest.importorskip("torch") +from families.phi4_multimodal import default_decoder # noqa: E402 from families.phi4_multimodal import model as family_model # noqa: E402 from families.phi4_multimodal.config import ModelConfig # noqa: E402 @@ -85,6 +86,8 @@ def _run(engine, cache_rows: int, caches: dict, tokens: list[int], position: int value = caches[name][:cache_rows] else: value = inputs[name].cuda() + if -1 not in tuple(engine.get_tensor_shape(name)): # the decode engine's static inputs + value = value.reshape(tuple(engine.get_tensor_shape(name))) bound[name] = value if -1 in tuple(engine.get_tensor_shape(name)): assert context.set_input_shape(name, tuple(value.shape)), name @@ -101,26 +104,18 @@ def _run(engine, cache_rows: int, caches: dict, tokens: list[int], position: int return {name: value.cpu().numpy() for name, value in outputs.items()} -@pytest.mark.gpu -@pytest.mark.trt -def test_prefill_over_the_cached_prefix_matches_the_whole_cache(tmp_path: Path) -> None: - if not torch.cuda.is_available(): - pytest.skip("requires CUDA") +def _engine(tmp_path: Path, role: str): config = _write_checkpoint(tmp_path) config.raw["_model_dir"] = str(tmp_path) - config.raw["_decoder_engine_role"] = "prefill" + config.raw["_decoder_engine_role"] = role instance = family_model._Phi4MultimodalModel() weights = instance.load_weights(str(tmp_path), config) - engine = trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine( + return trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine( instance.build_engine(config, weights, MAX_CACHE, precision="fp32") ) - assert tuple(engine.get_tensor_shape("cache_k_0"))[0] == -1 - assert [tuple(s) for s in engine.get_tensor_profile_shape("cache_k_0", 0)] == [ - (1, engine.get_tensor_shape("cache_k_0")[1]), - (1, engine.get_tensor_shape("cache_k_0")[1]), - (MAX_CACHE, engine.get_tensor_shape("cache_k_0")[1]), - ] + +def _random_caches(engine) -> dict: caches = {} for i in range(engine.num_io_tensors): name = engine.get_tensor_name(i) @@ -128,6 +123,23 @@ def test_prefill_over_the_cached_prefix_matches_the_whole_cache(tmp_path: Path) width = engine.get_tensor_shape(name)[1] # Rows beyond the prefix hold garbage that the mask must hide. caches[name] = torch.randn(MAX_CACHE, width, device="cuda") + return caches + + +@pytest.mark.gpu +@pytest.mark.trt +def test_prefill_over_the_cached_prefix_matches_the_whole_cache(tmp_path: Path) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + engine = _engine(tmp_path, "prefill") + assert tuple(engine.get_tensor_shape("cache_k_0"))[0] == -1 + assert [tuple(s) for s in engine.get_tensor_profile_shape("cache_k_0", 0)] == [ + (1, engine.get_tensor_shape("cache_k_0")[1]), + (1, engine.get_tensor_shape("cache_k_0")[1]), + (MAX_CACHE, engine.get_tensor_shape("cache_k_0")[1]), + ] + + caches = _random_caches(engine) rng = np.random.default_rng(0) for position, rows in ((0, 1), (5, 16)): @@ -136,3 +148,28 @@ def test_prefill_over_the_cached_prefix_matches_the_whole_cache(tmp_path: Path) prefix = _run(engine, rows, caches, tokens, position) for name, value in whole.items(): np.testing.assert_allclose(prefix[name], value, rtol=1e-4, atol=1e-4, err_msg=name) + + +@pytest.mark.gpu +@pytest.mark.trt +def test_decode_over_the_cached_prefix_matches_the_whole_cache(tmp_path: Path) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + engine = _engine(tmp_path, "decode") + width = engine.get_tensor_shape("cache_k_0")[1] + opt_rows = min(MAX_CACHE, default_decoder.DECODE_OPT_CACHE_ROWS) + assert tuple(engine.get_tensor_shape("cache_k_0"))[0] == -1 + assert [tuple(s) for s in engine.get_tensor_profile_shape("cache_k_0", 0)] == [ + (1, width), + (opt_rows, width), + (MAX_CACHE, width), + ] + caches = _random_caches(engine) + + rng = np.random.default_rng(1) + for position, rows in ((0, 1), (5, 16), (40, 48), (MAX_CACHE - 1, MAX_CACHE)): + token = rng.integers(0, VOCAB, 1).tolist() + whole = _run(engine, MAX_CACHE, caches, token, position) + prefix = _run(engine, rows, caches, token, position) + for name, value in whole.items(): + np.testing.assert_allclose(prefix[name], value, rtol=1e-4, atol=1e-4, err_msg=name)