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/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/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 0c98c8d9a0..7a3cf916d1 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; @@ -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(mask_width)) mask_buf_.resize(static_cast(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 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 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())) + 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); - 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 +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 { diff --git a/families/phi4_multimodal/runtime/kv_cache.h b/families/phi4_multimodal/runtime/kv_cache.h index 54d85289b2..297b14a3b1 100644 --- a/families/phi4_multimodal/runtime/kv_cache.h +++ b/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 cache_k_; // [num_layers], shape [max_length, kv_dim] std::vector 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 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..a786000627 --- /dev/null +++ b/families/phi4_multimodal/tests/test_prefill_cache_rows.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The prefill and decode engines attend 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 default_decoder # noqa: E402 +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() + 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 + 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()} + + +def _engine(tmp_path: Path, role: str): + config = _write_checkpoint(tmp_path) + config.raw["_model_dir"] = str(tmp_path) + config.raw["_decoder_engine_role"] = role + instance = family_model._Phi4MultimodalModel() + weights = instance.load_weights(str(tmp_path), config) + return trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine( + instance.build_engine(config, weights, MAX_CACHE, precision="fp32") + ) + + +def _random_caches(engine) -> dict: + 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") + 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)): + 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) + + +@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)