diff --git a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md index cf84dd03896f..fc710b8dd716 100644 --- a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md +++ b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md @@ -3,7 +3,7 @@ name: trtllm-model-onboard-multimodal description: > Onboard a HuggingFace multimodal model (vision/audio/video + text) to the TensorRT-LLM PyTorch backend. Use when writing a new - `tensorrt_llm/_torch/models/modeling_.py` plus its input processor and + `tensorrt_llm/_torch/models/modeling_{vlm}.py` plus its input processor and weight mapper, or extending an existing VLM. Not for AutoDeploy — use `ad-model-onboard` for that path. license: Apache-2.0 @@ -52,23 +52,26 @@ metadata: executor to drive KV-cache hash matching. [4] Per-iteration staging (model engine) - Context: build MultimodalRuntimeData (positions / lengths / chunk bounds) - → push pixel_values to CUDA pinned + non_blocking, obeying the - model's multimodal_data_device_paths declaration; pad - mrope_position_ids into a preallocated CUDA buffer. + Context: build MultimodalRuntimeData (positions / lengths / chunk bounds). + For item-scheduled models, reserve space in the model's one + TensorLRUCache. Reuse ready outputs and encode each missing item only + once, within the encoder item / token / output-byte limits. + Legacy models stage the request payload here as before. H2D obeys + multimodal_data_device_paths and uses pinned, non-blocking copies. Generation (mRoPE only): strip everything except mrope_position_deltas. - Post-prefill: drop mm_data so it doesn't ride along in decode. + Post-prefill / termination: release request-held cache references, then strip + raw MM fields so they don't ride along in decode. [5] Model.forward(attn_metadata, input_ids, position_ids, multimodal_params=…) - get_multimodal_embeddings: runs encoder.forward only on params whose - multimodal_data["multimodal_embedding"] is empty (chunked-prefill iter - 2+ hits the per-request cache; results written back automatically). - find_input_mm_embeds: slices the cached embedding to the current chunk - under chunked prefill / KV-cache reuse. + MultimodalModelMixin.prepare_multimodal_inputs: concatenates prompt-ordered + cache-owned item outputs once into the existing embedding-tensor contract; + the legacy path still uses get_multimodal_embeddings and its + full-request/partial-hit behavior. + find_input_mm_embeds: slices active chunk rows from that tensor. prepare_mrope_config (mRoPE models): one-shot mrope_rotary_cos_sin per request from the staged mrope_position_ids buffer. - fuse_input_embeds: text + mm merged via precomputed indices - (with optional extra_embeds for multi-feature encoders). + fuse_input_embeds: merges text and MM rows through the existing precomputed + index path (optional extra_embeds remain available for multi-feature encoders). self.llm.forward(inputs_embeds=..., mrope_config=...) → logits. ``` @@ -77,6 +80,14 @@ metadata: - The producer hands off as **handles** at the end of [2], so the broadcast in [3] stays small (Contract 3). - [4] is the only per-iteration GPU staging; H2D is `non_blocking=True` from pinned host memory. - [5] runs on the compute stream and must be sync-free (Contract 1). +- Item-scheduled encoder outputs live only in the model-owned `TensorLRUCache`. + Each request remembers the cache entry for every item, in prompt order, but + does not keep a second output tensor or combined output buffer. The cache is + large enough for one legal encoder iteration; `encoder_cache_max_bytes` may + make it larger when reuse is enabled. Stable item keys let requests share + one encoded output. Items without a stable key use a request-local key. The + single tensor assembled for an LLM forward is temporary and is not retained + as a second request-owned cache. ### EPD-disaggregated path @@ -85,9 +96,15 @@ When `@support_multimodal_disaggregated` is set and the deployment uses `TLLM_MU - **Encoder worker:** runs as a standalone `MultimodalEncoder` (`mm_encoder_only=True`). It executes only the multimodal encoder and ships `mm_embeddings` (+ mRoPE position ids/deltas) to prefill+decode workers as shared-tensor handles. - **Prefill+decode worker:** the model's `__init__` skips constructing `self.mm_encoder` when `_is_mm_disagg()` is true; the input processor's `attach_multimodal_embeddings()` override binds the encoder handles into the request. For context-only requests, the engine re-clones mrope tensors so IPC handles outlive the encoder worker's freed memory — replicate that pattern for any new GPU-resident mm tensors. +This item-scheduled cache path currently applies only when the encoder and LLM +run in the same process. `mm_encoder_only` / EPD keeps its existing +shared-handle transfer path. Item scheduling also rejects side-stream encoder +prefetch today, and LLM prefill waits until every item in a request is ready. +Side-stream support and item/LLM prefill overlap are separate follow-ups. + ### Templates to study -`modeling_qwen3vl.py`, `modeling_llava_next.py`, and `modeling_gemma3vl.py` are the canonical references — fully-ported encoder, single-class wrapper, `text_config`-based LLM resolution. Other examples by modality: `modeling_pixtral.py`, `modeling_phi4mm.py` (audio), `modeling_mllama.py`, `modeling_hyperclovax.py`, `modeling_mistral_large3.py`. Pick the closest one (modality + LLM family + RoPE variant). `modeling_qwen2vl.py` retains an HF-passthrough vision tower for the outdated Qwen2-VL family — read it for context but don't copy that pattern. +`modeling_qwen3vl.py` and `modeling_mistral.py` are the canonical references for `MultimodalModelMixin`, item scheduling, and the unified encoder-output cache. Use `modeling_llava_next.py` and `modeling_gemma3vl.py` for modality/family details while recognizing that their runtime path is not migrated yet. Other examples: `modeling_phi4mm.py` (audio), `modeling_mllama.py`, `modeling_hyperclovax.py`. `modeling_qwen2vl.py` retains an HF-passthrough vision tower for the outdated Qwen2-VL family; read it for context but don't copy that pattern. --- @@ -176,7 +193,11 @@ A 1024×1024 fp32 patch tensor is ~12 MB; a video clip can be hundreds of MB. Na - **Always use `MultimodalParams.to_handle`/`to_tensor`.** `to_handle` swaps each tensor inside `multimodal_data` for a small dict — `{method_key, tensor_size, storage_handle, ...}` — that points at the same memory: a CUDA-IPC handle for GPU tensors (`REBUILD_CUDA`) or a POSIX-shm handle for CPU tensors (`REBUILD_CPU`). The dict is a few hundred bytes regardless of the original tensor size. Consumers call `to_tensor` to rebuild local tensor views from the handle. See `_torch/shared_tensor/`. - **Where it crosses ranks:** the executor broadcasts `py_multimodal_data` via `dist.broadcast` / `tp_cp_broadcast` / PP send-recv. Payload size = the literal byte size of whatever's in `py_multimodal_data` — confirm every tensor inside has been swapped for its handle dict (i.e. `to_handle` ran) before this point. -- **Strip after prefill.** `_strip_py_multimodal_data_post_prefill` clears everything except `mrope_config.mrope_position_deltas`. If your model needs to retain something across decode, update `strip_mm_data_for_generation` explicitly. +- **Release and strip after prefill.** `PyExecutor._release_multimodal_resources` + releases the request's cache entries and then calls + `strip_mm_data_for_generation`. Repeated cleanup is safe. If your model needs + data during decode, update `strip_mm_data_for_generation`; do not add another + post-prefill strip helper. - **EPD disagg.** Embeddings still cross workers as shared tensors, not bytes — see the EPD-disaggregated path section above for the encoder/prefill-worker split. - **Hashes are small; broadcast eagerly.** `MultimodalInput.multimodal_hashes` (blake3) drives KV-cache reuse — never substitute raw pixels for them. @@ -184,9 +205,20 @@ A 1024×1024 fp32 patch tensor is ~12 MB; a video clip can be hundreds of MB. Na ### Contract 4 — Batch the multimodal encoder across requests -`get_multimodal_embeddings` hands the encoder a **list** of `MultimodalParams` covering every uncached request in the current batch. The encoder must consume that list as a single batched forward pass — concatenate every request's `pixel_values` / `image_grid_thw` / mel frames into one tensor, build one ad-hoc `attn_metadata` whose `seq_lens` carries per-image boundaries, and run the encoder blocks once. Looping `for p in mm_params: encoder.forward(p)` loses kernel-launch coalescing and serializes N requests' worth of encoder work. - -**Pattern (Qwen2.5-VL).** `Qwen2_5_VisionModel` concatenates every request's `pixel_values` into one `[total_patches, ...]` tensor and builds `attn_metadata` with `batch_size=1` and `seq_lens=[img1_patches, img2_patches, ...]`. The TRT-LLM `Attention` module respects `seq_lens` so cross-image attention doesn't bleed. The patch merger / projector at the end then splits the result back per-request via `torch.split` over the same lengths (this is what `_cache_multimodal_embeddings` expects too). +Implement `MultimodalModelMixin.encode_multimodal_inputs` as one batched forward over its +`MultimodalParams` list. The item-scheduled path calls `prepare_multimodal_encoder_inputs` for +only the selected producer items, then `forward_multimodal_encoder_items` groups compatible +modality runs and calls that same encoder hook. The legacy path also calls it for uncached +requests through `get_multimodal_embeddings`. Concatenate requests' `pixel_values` / +`image_grid_thw` / mel frames, build one ad-hoc `attn_metadata` whose `seq_lens` preserves item +boundaries, and run the encoder blocks once. Never loop over requests and call the encoder once +per request. + +**Pattern (Qwen3-VL).** `encode_multimodal_by_groups` concatenates items for compatible +modalities, runs one encoder call, and restores request/prompt order from `mm_item_order`. +`forward_multimodal_encoder_items` then splits the result into one tensor per atomic item using +the processor-declared output lengths. Override `build_multimodal_encoder_input` only when the +default packed-grid, stacked-image, or stacked-audio slicers cannot represent the model's input. **Audit.** Under load with several multimodal requests in one batch, the encoder kernels in nsys should appear as **one wide block per iteration**, not N narrow blocks. A fan of N narrow blocks means the encoder is being looped per request instead of batched — one of the easiest VLM perf regressions to introduce while refactoring. @@ -220,8 +252,11 @@ class {Name}VisionModel(nn.Module): ... -class {Name}Model(PreTrainedModel): +class {Name}Model(MultimodalModelMixin, PreTrainedModel): config_class = {Name}Config + supports_encoder_cache = True + # Enable only after implementing the Phase 3 atomic-item contracts. + supports_mm_encoder_item_scheduling = True def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs): config = model_config.pretrained_config @@ -248,41 +283,28 @@ class {Name}Model(PreTrainedModel): self.config = self.llm.config self.model_config.pretrained_config = self.llm.config - @property - def vocab_size_padded(self) -> int: - return self.llm.vocab_size_padded - - def infer_max_seq_len(self) -> int: - return self.llm.infer_max_seq_len() - - @torch.inference_mode() - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor] = None, - position_ids: Optional[torch.IntTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - return_context_logits: bool = False, - **kwargs, + def encode_multimodal_inputs( + self, multimodal_params: List[MultimodalParams] ) -> torch.Tensor: - num_context_requests = attn_metadata.num_contexts - - multimodal_params = kwargs.get("multimodal_params", []) - mm_embeds = [] - if len(multimodal_params) > 0 and not _is_mm_disagg(): - mm_embeds = get_multimodal_embeddings( - encoder_forward_fn=self.mm_encoder.forward, - multimodal_params=multimodal_params[:num_context_requests], - ) - mm_embeds = find_input_mm_embeds( - mm_embeds, multimodal_params[:num_context_requests]) - - input_ids, inputs_embeds = fuse_input_embeds( - self.llm.model.embed_tokens, input_ids, mm_embeds, **kwargs) - return self.llm.forward( - attn_metadata=attn_metadata, input_ids=input_ids, - position_ids=position_ids, inputs_embeds=inputs_embeds, - return_context_logits=return_context_logits) + if self.mm_encoder is None: + raise ValueError("Raw multimodal inputs require a local encoder") + return self.mm_encoder.forward(multimodal_params) + + @property + def language_model(self) -> torch.nn.Module: + return self.llm + + @property + def text_embedding_layer(self): + return self.llm.model.embed_tokens + + @property + def embedding_dim(self) -> int: + return self.text_embedding_layer.embedding_dim + + @property + def embedding_dtype(self) -> torch.dtype: + return self.text_embedding_layer.weight.dtype @property def multimodal_data_device_paths(self) -> List[str]: @@ -291,13 +313,26 @@ class {Name}Model(PreTrainedModel): **Required (every multimodal model):** -- `forward` takes `multimodal_params` via `**kwargs`. **Never** add `pixel_values` / `image_grid_thw` / `attention_mask` as direct args — they live in `multimodal_params.multimodal_data`. -- **Encoder output length must match the input processor's MM placeholder count.** `mm_encoder.forward` must return a single tensor whose first dimension equals the total number of MM tokens (excluding special tokens) the input processor placed in `prompt_token_ids`. If lengths don't agree — or if the encoder returns a list with more than one element — `get_multimodal_embeddings` silently skips caching the embedding back into `multimodal_data`, and chunked prefill re-runs the encoder from scratch on every chunk. +- Inherit `MultimodalModelMixin` and use its `forward` unless the family needs explicit + language-model delegation. If overriding, call `prepare_multimodal_inputs`; never add + `pixel_values` / `image_grid_thw` / `attention_mask` as direct args. +- Implement `encode_multimodal_inputs`, `language_model`, `text_embedding_layer`, + `embedding_dim`, and `embedding_dtype`. The encoder must return one tensor whose first + dimension equals the processor-declared aggregate output length. The item path validates and + splits it per item; the legacy path retains its existing full-request behavior. +- Set `supports_mm_encoder_item_scheduling=True` only when the input processor emits valid atomic + item metadata and the model can slice and batch selected items. Set `supports_encoder_cache=True` + only when the production forward consumes the mixin cache path. Do not construct or mutate a + request-local encoder-output cache in the model. **Family-specific extras (apply only when relevant):** -- **mRoPE (Qwen-VL family):** add `init_mrope_embedding(model_config)` in `__init__` to preallocate `self.mrope_position_ids_padding_cuda`, plus `prepare_mrope_config(multimodal_params, num_context_requests)` returning `mrope_rotary_cos_sin`. Pass through to `self.llm.forward(..., mrope_config=...)`. Reference: `Qwen3VLModelBase.prepare_mrope_config`. -- **Deepstack features (Qwen3-VL):** split encoder output into `mm_embed` + `deepstack_embeds`, call `fuse_input_embeds(..., extra_embeds=deepstack_embeds)`, forward `deepstack_embeds=` into the LLM. +- **mRoPE (Qwen-VL family):** add `init_mrope_embedding(model_config)` in `__init__`, then + build and forward the mRoPE config from `get_language_model_extra_forward_kwargs`. Reference: + `Qwen3VLModelBase.prepare_mrope_config`. +- **Deepstack features (Qwen3-VL):** keep each item's primary and deepstack rows in the same cache + entry, then use the existing `after_active_multimodal_embeddings` and + `_fuse_multimodal_embeddings` hooks after the prompt-order tensor is assembled. - **HF wrapper without a clean `text_config`:** Qwen2-VL's `Qwen2VLModelBase` rewrites `architectures` to surface the inner LLM. Fall back to that pattern only when the multimodal HF config does not expose a `text_config` sub-config. - **Inner LLM that doesn't match HF's `text_config` schema (Qwen3.5-MoE-VL → Qwen3Next).** When the VLM's HF `text_config` schema differs from the TRT-LLM runtime model you want to reuse, write a config normalizer (e.g. `_normalize_qwen35_moe_vl_config`) that maps HF aliases to the runtime's expected names (mRoPE keys, `intermediate_size` aliases, quantization-exclude module paths). Wire it via **lazy import** from `pyexecutor.config_utils.load_pretrained_config` — the `Mistral` and `Qwen3_5` branches are templates. Two gotchas: transformers 5.x's `rope_scaling` is a **property aliasing `rope_parameters`** — setting either silently overwrites the other, so the normalizer should mutate `rope_parameters` directly if the HF code still reads from it. And for VLMs, the normalizer must run on the **composite** config (with `text_config` / `vision_config`), not flattened away. - **Thin wrapper for runtime reuse.** Even when the LM class body is identical to the runtime's existing class, still create a `@register_auto_model("YourArch")`-decorated thin subclass — that's how weight-mapper dispatch picks the family-specific mapper. You can't stack two `@register_auto_model` decorators on a single shared class. @@ -317,6 +352,13 @@ The workspace dimension comes from `encoder_max_num_tokens`, falling back to the > Mixed image+video+audio models profile the supported modality with the largest legal per-item workload under the configured limits. Runtime mixed-modality requests still share the same aggregate token limit. +**Atomic-item scheduling.** Override `get_mm_encoder_item_metadata` to return +prompt-ordered `item_refs`, the encoder-token cost of each item, and each +item's output length. Also implement `get_max_mm_encoder_output_embeddings` so +the engine can size the cache for one legal encoder iteration. The declared +output lengths must match both the MM placeholder spans and the tensor splits +from `forward_multimodal_encoder_items`. + Implement `call_with_text_prompt(inputs, sampling_params)` — the per-model text-prompt path. **Don't override `__call__`**: the base class's concrete `__call__` dispatches here for text prompts, and also detokenizes `prompt_token_ids → prompt` and falls through to here for non-fast-path VLMs. `call_with_text_prompt` does: 1. Pull `text_prompt`, `mm_data`, `mm_processor_kwargs` from `inputs`. @@ -431,7 +473,7 @@ Follow `CONTRIBUTING.md`. Title `[JIRA/NVBUG/None][type] description`, `git comm **Architecture & registration** - [ ] Decorator stack in correct order: `@support_multimodal_disaggregated` (outermost, optional) → `@register_vision_encoder` → `@register_auto_model` → `@register_input_processor` (innermost). -- [ ] `forward` takes `multimodal_params` via `**kwargs`; no `pixel_values` / `image_grid_thw` / `attention_mask` direct args. +- [ ] Model inherits `MultimodalModelMixin`; an overridden `forward` calls `prepare_multimodal_inputs` and takes no raw MM tensor args. - [ ] `multimodal_data_device_paths` lists every GPU-resident mm field. - [ ] If runtime-reusing (e.g. Qwen3.5 → Qwen3Next): thin `@register_auto_model` wrapper class present; config normalizer lazy-imported from `pyexecutor.config_utils.load_pretrained_config`. @@ -444,7 +486,7 @@ Follow `CONTRIBUTING.md`. Title `[JIRA/NVBUG/None][type] description`, `git comm **Input processor** - [ ] Subclasses both `BaseMultimodalInputProcessor` and `BaseMultimodalDummyInputsBuilder`. -- [ ] Encoder KV-cache profiling: implements the deterministic dummy contract (`get_mm_max_tokens_per_item` + `get_dummy_mm_data`) and the model exposes `encode_multimodal_inputs`; encoder inherits `MultimodalEncoderMixin` (no hardcoded `max_num_*=8192` — sized by `setup_attn_metadata`). Skipping these = text-only dummy, encoder memory unaccounted. +- [ ] Encoder profiling implements `get_mm_max_tokens_per_item` + `get_dummy_mm_data`; item scheduling additionally implements `get_mm_encoder_item_metadata` + `get_max_mm_encoder_output_embeddings`. The model exposes batched `encode_multimodal_inputs`; encoder metadata capacity comes from `setup_attn_metadata`, not hardcoded `max_num_*` values. - [ ] `call_with_text_prompt` (not `__call__` — that's the base-class dispatcher) runs HF AutoProcessor + tokenizer, builds `multimodal_data` by modality, computes `mrope_config` on CPU, `_postprocess`-rewrites mm token ids to the OOV sentinel. - [ ] `mm_processor_kwargs` flow-through preserved. (Tokenized fast path is optional: set `supports_token_id_mm_expansion = True` + implement `get_text_with_mm_placeholders` / `expand_prompt_token_ids_for_mm`; otherwise the base class detokenizes token-ID inputs automatically.) - [ ] `_attach_multimodal_embeddings_impl` implemented (not the `attach_multimodal_embeddings` wrapper) if `@support_multimodal_disaggregated`. @@ -454,8 +496,8 @@ Follow `CONTRIBUTING.md`. Title `[JIRA/NVBUG/None][type] description`, `git comm - [ ] `set_sync_debug_mode("warn")` audit on prefill: zero warnings from your model. - [ ] Async loaders used for URL/bytes inputs. - [ ] Broadcast payload < 1 MB per rank per request (NVTX `broadcast_requests` / `tp_broadcast_requests`); media crosses ranks only via `to_handle` / `to_tensor`. -- [ ] Decode-iteration `mm_data` is empty (post-prefill strip exercised in e2e test). -- [ ] Encoder output is a single tensor whose first dim equals the input processor's MM placeholder count; verified by running with chunked prefill on (small `--max_num_tokens`) and confirming the encoder runs once per request, not once per chunk. +- [ ] Post-prefill/termination cleanup drains item-cache refs exactly once and leaves only fields retained by `strip_mm_data_for_generation`. +- [ ] Encoder output is one tensor whose first dim equals the processor-declared MM output rows; cache item splits preserve prompt order and are reassembled into the existing single-tensor fusion contract. - [ ] Encoder is batched across requests: a multi-request batch produces a single wide encoder block in nsys, not N narrow blocks (Contract 4). **Tests & docs** diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index f4046b208bf6..cd0f44afe0b6 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -1066,7 +1066,7 @@ def _get_or_encode_multimodal_embeddings( records each entry's producer event on the issuing (aux) stream; the next iteration's main-stream consumer waits on the request-level `encoder_event` for ordering. """ - encoder_cache = self._get_multimodal_encoder_cache() + encoder_cache = self._multimodal_encoder_cache cache_misses: list[MultimodalParams] = [] partial_hits: list[tuple[MultimodalParams, EncoderCachePartition]] = [] if encoder_cache is not None: @@ -1107,60 +1107,51 @@ def _get_or_encode_multimodal_embeddings( self._validate_embeddings(embeddings, multimodal_params) return embeddings[0] - def _get_multimodal_encoder_cache(self) -> Optional[TensorLRUCache]: - """Return the per-model full-request-path encoder clone cache, if enabled. + def _initialize_multimodal_encoder_cache(self, max_bytes: int) -> Optional[TensorLRUCache]: + """Initialize the model-owned multimodal encoder-output cache once. The cache stores per-item embeddings for params that can be represented by one modality. See `_encoder_cache_keys` for the mixed-modality skip path and its technical limitation. - - Scope: the single encoder cache instance for a cache-enabled model - (`supports_encoder_cache`). The full-request (legacy inline-encode) consumers — side-stream - prefetch, `mm_encoder_only`/disagg encoding — populate and read it inline; the - item-scheduling path consumes the same instance read-through at encode time - (`ModelEngine.forward_multimodal_encoder_items`). The key format is shared - (`_encoder_cache_item_key`) so hits cross between paths. The item path's recorded outputs - are cloned, so cache eviction never invalidates an in-flight request. + `ModelEngine` resolves `max_bytes` as the larger of the item-scheduling + output budget and persistent-reuse capacity before runtime access. + Zero leaves the cache disabled. """ - if not self.encoder_cache_active: + if self._multimodal_encoder_cache is not None: + return self._multimodal_encoder_cache + + if max_bytes == 0: logger.debug_once( - f"{_MM_ENCODER_CACHE_LOG_NAME}: disabled because the model does not opt in via " - "supports_encoder_cache or multimodal_config.encoder_cache_max_bytes=0.", + f"{_MM_ENCODER_CACHE_LOG_NAME}: disabled because neither item scheduling nor " + "persistent reuse requires storage.", key="mm_encoder_cache_disabled", ) return None multimodal_config = self.model_config.multimodal_config - max_bytes = multimodal_config.encoder_cache_max_bytes - if self._multimodal_encoder_cache is None: - # Per-item embeddings are views produced by splitting a request-level encoder output. - # Clone them so a cached item neither aliases mutable caller output nor retains the - # entire batch allocation while cache accounting charges only its logical size. This - # briefly needs source and clone memory during insertion, but preserves existing cache - # entries when the copy cannot be allocated. - self._multimodal_encoder_cache = TensorLRUCache( - max_bytes, - name=_MM_ENCODER_CACHE_LOG_NAME, - cuda_stream_aware=multimodal_config.encoder_side_stream_max_ahead > 0, + self._multimodal_encoder_cache = TensorLRUCache( + max_bytes, + name=_MM_ENCODER_CACHE_LOG_NAME, + cuda_stream_aware=multimodal_config.encoder_side_stream_max_ahead > 0, + ) + try: + embedding_dim = self.embedding_dim + embedding_dtype = self.embedding_dtype + except NotImplementedError: + logger.info( + f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, " + "embedding row capacity unavailable because the model does not implement " + "embedding_dim and embedding_dtype." + ) + else: + bytes_per_embedding_row = ( + embedding_dim * torch.empty((), dtype=embedding_dtype).element_size() + ) + max_embedding_rows = max_bytes // bytes_per_embedding_row + logger.info( + f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, " + f"max_embedding_rows={max_embedding_rows}, embedding_dim={embedding_dim}, " + f"embedding_dtype={embedding_dtype}" ) - try: - embedding_dim = self.embedding_dim - embedding_dtype = self.embedding_dtype - except NotImplementedError: - logger.info( - f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, " - "embedding row capacity unavailable because the model does not implement " - "embedding_dim and embedding_dtype." - ) - else: - bytes_per_embedding_row = ( - embedding_dim * torch.empty((), dtype=embedding_dtype).element_size() - ) - max_embedding_rows = max_bytes // bytes_per_embedding_row - logger.info( - f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, " - f"max_embedding_rows={max_embedding_rows}, embedding_dim={embedding_dim}, " - f"embedding_dtype={embedding_dtype}" - ) return self._multimodal_encoder_cache @staticmethod @@ -1678,7 +1669,7 @@ def _dispatch_cross_iter_prefetch( encoder_event = None try: with _run_on_aux_stream(aux_stream) as encoder_event: - encoder_cache = model._get_multimodal_encoder_cache() if encoder_cache_enabled else None + encoder_cache = model._multimodal_encoder_cache if encoder_cache_enabled else None cache_misses: list[MultimodalParams] = [] partial_hits: list[tuple[MultimodalParams, EncoderCachePartition]] = [] if encoder_cache is None: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index cd56117042ff..2d27615513bf 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -955,7 +955,7 @@ def _create_dummy_encoder_inputs(self) -> List[MultimodalParams]: return [MultimodalParams(multimodal_data=mm_data)] def _encode_dummy_inputs(self) -> Optional[torch.Tensor]: - """Run the full-budget MM encoder and retain request-owned output storage.""" + """Run the full-budget MM encoder and retain equivalent store output.""" if not self._dummy_encoder_inputs: return None @@ -975,7 +975,7 @@ def _encode_dummy_inputs(self) -> Optional[torch.Tensor]: ) output = self._model_engine.model.encode_multimodal_inputs( encoder_inputs) - # Runtime item state owns detached copies rather than views of + # The runtime store owns detached copies rather than views of # an encoder batch. Reproduce that allocation boundary here. return output.detach().clone() finally: @@ -984,20 +984,11 @@ def _encode_dummy_inputs(self) -> Optional[torch.Tensor]: def _get_multimodal_encoder_memory_reserve(self, profiled_output_bytes: int = 0 ) -> int: - """Return output and cache capacity absent from the measured peak.""" - output_budget = getattr(self._model_engine, - "mm_encoder_output_budget_bytes", None) - unprofiled_output_bytes = max(0, (output_budget or 0) - - profiled_output_bytes) - - model = self._model_engine.model - cache_bytes = 0 - if (isinstance(model, MultimodalModelMixin) - and model.encoder_cache_active - and model.model_config.multimodal_config is not None): - cache_bytes = ( - model.model_config.multimodal_config.encoder_cache_max_bytes) - return unprofiled_output_bytes + cache_bytes + """Return unified encoder-store capacity absent from the measured peak.""" + encoder_cache = self._model_engine.mm_encoder_cache + if encoder_cache is None: + return 0 + return max(0, encoder_cache.max_bytes - profiled_output_bytes) def _get_token_num_for_estimation(self) -> int: """Compute KV cache capacity required for estimate_max_kv_cache_tokens to succeed.""" @@ -1221,7 +1212,7 @@ def configure_kv_cache_capacity(self, if py_executor is not None and not self._skip_est: # Run the MM encoder at its independent token budget, then keep the - # resulting request-owned embeddings resident while the text-only + # equivalent cache-owned embeddings resident while the text-only # LLM dummy fills max_num_tokens. encoder_profile_output = self._encode_dummy_inputs() if encoder_profile_output is not None: @@ -3384,7 +3375,8 @@ def create_py_executor_instance( reorder_policy_config.policy_args.agent_inflight_seq_num) scheduler = SimpleScheduler(capacity_scheduler, mb_scheduler) - if getattr(model_engine, "mm_encoder_item_scheduling_enabled", False): + if (getattr(model_engine, "mm_encoder_item_scheduling_enabled", False) + and model_engine.mapping.is_first_pp_rank()): # Wrap the LLM scheduler with atomic MM item budgeting. ModelEngine # already validated model-capability-dependent feature combinations. multimodal_config = llm_args.multimodal_config @@ -3397,13 +3389,19 @@ def create_py_executor_instance( logger.info("Eager multimodal encoder scheduling is enabled for " "capacity-rejected active requests.") scheduler_cls = MultimodalEagerEncoderScheduler + encoder_cache = model_engine.mm_encoder_cache + if encoder_cache is None: + raise RuntimeError( + "MM encoder item scheduling requires its unified output cache") scheduler = scheduler_cls( scheduler, max_batch_size=model_engine.encoder_batch_size, max_num_tokens=model_engine.encoder_max_num_tokens, - output_budget_bytes=model_engine.mm_encoder_output_budget_bytes, + encoder_cache=encoder_cache, + get_item_cache_keys=model_engine.get_mm_encoder_item_cache_keys, bytes_per_encoder_embedding=( model_engine.bytes_per_mm_encoder_embedding), + retain_cache_entries=model_engine.model.encoder_cache_active, ) config = model_engine.model.model_config.pretrained_config diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py index 9effb6644d30..43d80acfd984 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py @@ -19,7 +19,7 @@ ) from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm._utils import prefer_pinned -from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.inputs.multimodal import MultimodalParams, strip_mm_encoder_inputs from tensorrt_llm.inputs.registry import ( BaseMultimodalDummyInputsBuilder, BaseMultimodalInputProcessor, @@ -29,7 +29,12 @@ from tensorrt_llm.llmapi.llm_args import MultimodalEncoderSchedulingPolicy, TorchLlmArgs from tensorrt_llm.logger import logger -from ..llm_request import LlmRequest, MultimodalEncoderRequestError, _Unset +from ..llm_request import ( + LlmRequest, + MultimodalEncoderProgress, + MultimodalEncoderRequestError, + _Unset, +) def resolve_mm_encoder_token_budget(base_budget: int, model_max_atomic_item_tokens: int) -> int: @@ -44,6 +49,12 @@ def validate_mm_encoder_scheduling_compatibility( policy = llm_args.multimodal_config.encoder_scheduling_policy if not item_scheduling_enabled: return + if llm_args.pipeline_parallel_size > 1: + raise ValueError( + "MM encoder item scheduling does not yet support pipeline " + "parallelism; set pipeline_parallel_size=1 or " + "encoder_scheduling_policy=DISABLED" + ) if llm_args.multimodal_config.encoder_side_stream_max_ahead > 0: raise ValueError( "MM encoder item scheduling does not yet support side-stream " @@ -155,14 +166,21 @@ def resolve_bytes_per_mm_encoder_embedding(model: MultimodalModelMixin) -> int: loaded weights' dtype -- both mixin properties are optional and most VLMs implement neither. """ + embedding_dim = None try: - return model.embedding_dim * torch.empty((), dtype=model.embedding_dtype).element_size() - except NotImplementedError: + embedding_dim = model.embedding_dim + except (AttributeError, NotImplementedError): pass + if embedding_dim is not None: + try: + embedding_dtype = model.embedding_dtype + except (AttributeError, NotImplementedError): + embedding_dtype = model.model_config.torch_dtype + return embedding_dim * torch.empty((), dtype=embedding_dtype).element_size() try: weight = model.text_embedding_layer.weight return weight.shape[-1] * weight.element_size() - except NotImplementedError: + except (AttributeError, NotImplementedError): pass pretrained = model.model_config.pretrained_config hidden_size = getattr(pretrained, "hidden_size", None) @@ -175,7 +193,7 @@ def resolve_bytes_per_mm_encoder_embedding(model: MultimodalModelMixin) -> int: "text_embedding_layer, and its pretrained config exposes " "no (text_config.)hidden_size" ) - element_size = next(model.parameters()).dtype.itemsize + element_size = model.model_config.torch_dtype.itemsize return hidden_size * element_size @@ -193,9 +211,9 @@ def resolve_mm_encoder_output_budget( that embedding capacity multiplied by bytes per encoder embedding; the LLM-side ``max_num_tokens`` does not participate. - It caps outputs held between encode and prefill and is reserved during KV-capacity - estimation. It is separate from the optional reuse cache (``encoder_cache_max_bytes``). - A request whose total embedding exceeds this budget is rejected at admission. + It is the minimum capacity of the unified encoder-output cache and is reserved during + KV-capacity estimation. Optional reuse may make the same cache larger. A request whose + total embedding exceeds this budget is rejected at admission. The embedding capacity is validated before the model is consulted, so a processor that cannot report one raises regardless of what the model implements. @@ -214,7 +232,7 @@ def resolve_mm_encoder_output_budget( class MultimodalItemScheduler: - """Encodes scheduler-selected MM items through an optional read-through cache. + """Encodes scheduler-selected MM items through the unified output cache. Constructed once, at engine startup, and only when item scheduling is engaged. It holds no reference to the engine: everything it reads is passed in. @@ -300,9 +318,8 @@ def create( embeddings produced by one legal encoder iteration, converted to bytes. Enforced by the scheduler; any capacity not materialized by warmup is reserved in KV-capacity estimation. - * (D) reuse cache bytes -- ``encoder_cache_max_bytes`` on the mixin's - ``TensorLRUCache``, self-bounded by LRU; unprofiled capacity is reserved on top - of (C) for cache-enabled models. + * (D) reuse cache bytes -- ``encoder_cache_max_bytes`` may make the same + ``TensorLRUCache`` larger than (C); it does not create a second pool. Prefill currently waits for every item in a request, so admission rejects a request whose complete MM embedding exceeds (C). @@ -371,37 +388,16 @@ def create( @property def encoder_cache(self) -> TensorLRUCache[Any] | None: - """The encoder cache the item path reads through, or ``None``. - - Only models that opt into the encoder cache (``supports_encoder_cache`` with - ``encoder_cache_max_bytes > 0``) participate -- the same lazily created - ``TensorLRUCache`` instance the legacy inline path uses. Item-scheduled models - without the flag (e.g. Qwen today) get ``None`` and never touch the cache; - cross-request reuse for them is out of scope here. - """ - model = self.model - if not getattr(model, "supports_encoder_cache", False): - return None - getter = getattr(model, "_get_multimodal_encoder_cache", None) - return getter() if getter is not None else None + """The one model-owned encoder-output cache used by item scheduling.""" + return self.model._multimodal_encoder_cache - def item_keys(self, request: LlmRequest) -> list[Hashable] | None: - """Return the request's per-item cache keys, or ``None``. - - ``None`` means the request cannot build stable content keys (no item metadata, - content hashes, or processor-kwargs hash), so its items bypass the cache (always - encoded; outputs live only on the request). The key format is shared with the - legacy full-request path (``_encoder_cache_item_key``) so entries hit across both. - - The engine only builds this component when item scheduling is engaged, so there is - no "is scheduling on" guard here. - """ - # Memoized on the request's encoder state: the inputs are fixed at - # admission, and a request whose items span several iterations would - # otherwise rebuild the same keys on each one. + def item_cache_keys(self, request: LlmRequest) -> list[Hashable] | None: + """Return stable per-item cache keys, or ``None`` for request-local keys.""" + # Request inputs do not change after admission. Cache these keys so a + # multi-iteration request does not rebuild them every time. state = request.py_mm_encoder_state - if state is not None and not isinstance(state.cache_item_keys, _Unset): - return state.cache_item_keys + if state is not None and not isinstance(state.stable_item_cache_keys, _Unset): + return state.stable_item_cache_keys mm_data = request.py_multimodal_data try: item_metadata = get_multimodal_encoder_item_metadata(mm_data) @@ -420,7 +416,7 @@ def item_keys(self, request: LlmRequest) -> list[Hashable] | None: ) ) if state is not None: - state.cache_item_keys = keys + state.stable_item_cache_keys = keys return keys @torch.inference_mode() @@ -429,96 +425,126 @@ def forward_items( requests: list[LlmRequest], scheduled_items: dict[int, list[int]], ) -> None: - """Forward selected MM encoder items and commit request-local outputs.""" + """Encode selected producer items into their reserved cache entries.""" if not scheduled_items: return if not isinstance(self.model, MultimodalModelMixin): raise TypeError("Item-level MM scheduling requires MultimodalModelMixin") - # Read-through against the model's encoder cache when enabled - # (`supports_encoder_cache` + `encoder_cache_max_bytes > 0`): a hit - # records a clone and skips the encode; a miss encodes, records, and - # populates the cache. The hit/miss branch runs here -- at encode time, - # on every rank against rank-local cache state -- so ranks perform - # identical get/put sequences and stay in sync. When the cache is off - # (`encoder_cache` is None), every item is a miss and outputs live only - # on the request. Records take a clone regardless, so a recorded slot - # never aliases a batch output or an evictable cache entry. encoder_cache = self.encoder_cache + if encoder_cache is None: + raise RuntimeError("MM item scheduling requires a model-owned encoder cache") request_by_id = {request.request_id: request for request in requests} - miss_items = [] - miss_owners: list[tuple[LlmRequest, int, Hashable | None]] = [] - touched_requests: dict[int, LlmRequest] = {} + encoder_items = [] + output_targets: list[tuple[int, Hashable, int]] = [] + scheduled_cache_keys: set[Hashable] = set() + + def requests_using_cache_keys(cache_keys: set[Hashable]) -> set[int]: + return { + request.request_id + for request in requests + if request.py_mm_encoder_state is not None + and any( + cache_key in cache_keys + for cache_key in request.py_mm_encoder_state.item_cache_keys + if cache_key is not None + ) + } + for request_id, item_indices in scheduled_items.items(): request = request_by_id.get(request_id) if request is None: raise MultimodalEncoderRequestError( - f"Scheduled MM request {request_id} is no longer active" + f"Scheduled MM request {request_id} is no longer active", + request_ids={request_id}, ) state = request.py_mm_encoder_state if state is None: raise MultimodalEncoderRequestError( - f"Scheduled MM request {request_id} has no encoder item state" + f"Scheduled MM request {request_id} has no encoder item state", + request_ids={request_id}, ) - touched_requests[request_id] = request multimodal_param = MultimodalParams(multimodal_data=request.py_multimodal_data) - # Scope the lookup to this iteration's items: probing an item the - # budget cannot encode yet would still refresh its LRU recency and - # reorder eviction against items actually in flight. - # Keys come from the request: the executor's params carry - # `py_multimodal_data` only, while the content hashes live on the - # `LlmRequest`. - item_keys = self.item_keys(request) if encoder_cache is not None else None - try: - partition = ( - self.model.partition_encoder_cache( - multimodal_param, encoder_cache, item_indices=item_indices, keys=item_keys + for item_idx in item_indices: + cache_key = state.item_cache_keys[item_idx] + if cache_key is None: + raise MultimodalEncoderRequestError( + f"Scheduled MM item {item_idx} has no cache key", + request_ids={request_id}, ) - if item_keys is not None - else None - ) - except MultimodalEncoderContractError as error: - raise MultimodalEncoderRequestError(str(error)) from error - if partition is None: - # No cache, or the request cannot build stable content keys: - # every scheduled item is a miss and its output lives only on - # the request. - for item_idx in item_indices: - miss_items.append((multimodal_param, item_idx)) - miss_owners.append((request, item_idx, None)) - continue - for item_idx, cached in partition.hits.items(): - state.record(item_idx, cached) - for item_idx in partition.miss_indices: - miss_items.append((multimodal_param, item_idx)) - miss_owners.append((request, item_idx, partition.keys[item_idx])) - - if miss_items: - try: - encoder_inputs = self.model.prepare_multimodal_encoder_inputs(miss_items) - except MultimodalEncoderContractError as error: - raise MultimodalEncoderRequestError(str(error)) from error - for encoder_input, _, _ in encoder_inputs: - encoder_input.to_device( - "multimodal_data", - "cuda", - pin_memory=prefer_pinned(), - target_keywords=getattr(self.model, "multimodal_data_device_paths", None), - ) + scheduled_cache_keys.add(cache_key) + encoder_items.append((multimodal_param, item_idx)) + output_targets.append((item_idx, cache_key, state.embedding_lengths[item_idx])) - try: - outputs = self.model.forward_multimodal_encoder_items(encoder_inputs) - except MultimodalEncoderContractError as error: - raise MultimodalEncoderRequestError(str(error)) from error - if len(outputs) != len(miss_owners): + try: + encoder_inputs = self.model.prepare_multimodal_encoder_inputs(encoder_items) + except MultimodalEncoderContractError as error: + raise MultimodalEncoderRequestError( + str(error), request_ids=requests_using_cache_keys(scheduled_cache_keys) + ) from error + for encoder_input, _, _ in encoder_inputs: + encoder_input.to_device( + "multimodal_data", + "cuda", + pin_memory=prefer_pinned(), + target_keywords=getattr(self.model, "multimodal_data_device_paths", None), + ) + + try: + outputs = self.model.forward_multimodal_encoder_items(encoder_inputs) + except MultimodalEncoderContractError as error: + raise MultimodalEncoderRequestError( + str(error), request_ids=requests_using_cache_keys(scheduled_cache_keys) + ) from error + if len(outputs) != len(output_targets): + raise MultimodalEncoderRequestError( + "MM item encoder must return one output per item", + request_ids=requests_using_cache_keys(scheduled_cache_keys), + ) + + for output, (item_idx, cache_key, expected_rows) in zip( + outputs, output_targets, strict=True + ): + if output.shape[0] != expected_rows: raise MultimodalEncoderRequestError( - "MM item encoder must return one output per item" + f"MM item {item_idx} produced {output.shape[0]} embeddings; " + f"expected {expected_rows}", + request_ids=requests_using_cache_keys({cache_key}), ) + encoder_cache.commit(cache_key, output) + for live_request in requests: + live_state = live_request.py_mm_encoder_state + if live_state is not None: + live_state.mark_cache_key_ready(cache_key) + + for request in requests: + state = request.py_mm_encoder_state + if state is not None and state.progress is MultimodalEncoderProgress.READY: + strip_mm_encoder_inputs(request.py_multimodal_data) - for output, (request, item_idx, key) in zip(outputs, miss_owners, strict=True): - request.py_mm_encoder_state.record(item_idx, output) - if encoder_cache is not None and key is not None: - encoder_cache.put(key, output) + def build_multimodal_data_for_llm(self, request: LlmRequest) -> dict[str, Any] | None: + """Attach prompt-ordered cached item outputs for LLM prefill.""" + state = request.py_mm_encoder_state + if state is None: + return request.py_multimodal_data + if state.progress is not MultimodalEncoderProgress.READY: + raise MultimodalEncoderRequestError( + f"MM request {request.request_id} reached prefill before its encoder outputs " + "were ready" + ) + encoder_cache = self.encoder_cache + if encoder_cache is None: + raise RuntimeError("MM request state requires an encoder cache") + + segments: list[torch.Tensor] = [] + for item_idx, cache_key in enumerate(state.item_cache_keys): + segment = encoder_cache.get(cache_key, record_stats=False) + if segment is None: + raise MultimodalEncoderRequestError( + f"Ready MM item {item_idx} is absent from the encoder cache" + ) + segments.append(segment) - for request in touched_requests.values(): - request.py_mm_encoder_state.finalize(request.py_multimodal_data) + multimodal_data = dict(request.py_multimodal_data or {}) + multimodal_data["multimodal_embedding"] = torch.cat(segments, dim=0) + return multimodal_data diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 02c01a2cfc8a..5e54e7b44740 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -2,12 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """Python extensions for executor requests.""" -import itertools from copy import copy, deepcopy from dataclasses import dataclass, field from enum import Enum, auto -from typing import (TYPE_CHECKING, Any, Dict, Hashable, List, Optional, Union, - cast) +from typing import (TYPE_CHECKING, Any, Dict, Hashable, Iterable, List, + Optional, Union, cast) import torch @@ -16,7 +15,6 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.bindings import executor as tllm_executor from tensorrt_llm.executor.result import SimpleTokenLogprobs, TokenLogprobs -from tensorrt_llm.inputs.multimodal import strip_mm_encoder_inputs from tensorrt_llm.inputs.registry import get_multimodal_encoder_item_metadata from tensorrt_llm.sampling_params import LogprobMode @@ -58,6 +56,14 @@ class MultimodalEncoderRequestError(ValueError): """A request-scoped MM encoder state or output contract violation.""" + def __init__(self, + message: str, + *, + request_ids: Optional[Iterable[int]] = None) -> None: + super().__init__(message) + # A shared cache-entry failure may need to fail multiple requests. + self.request_ids = frozenset(request_ids or ()) + class MultimodalEncoderProgress(Enum): """Python-only progress derived from request-local MM item outputs. @@ -90,39 +96,32 @@ class _Unset: _UNSET = _Unset() +def make_mm_encoder_transient_cache_key(request_id: int, + item_idx: int) -> tuple[str, int, int]: + """Return a request-local cache key for an item without a stable key.""" + return ("mm_transient", request_id, item_idx) + + @dataclass class MultimodalEncoderRequestState: - """Per-item MM encoder bookkeeping owned by one request. - - Created at admission by `initialize_multimodal_encoder_request` for - requests whose raw MM payloads run through the item scheduler; the - request's `py_mm_encoder_state` is `None` otherwise. Items are written by - the single `record()` writer from two sources — fresh encoder outputs and - read-through encoder-cache hits — so validation cannot diverge between - them. - - The request's embeddings live in **one contiguous buffer**, sized from the - declared item lengths and allocated by the first `record()` (the point at - which the encoder's row shape, dtype and device become known). Each item - is copied into its own row range, so the buffer is the request's *final* - storage: it neither aliases the encoder's batch output nor a cache entry - that could be evicted under it, and the prefill path consumes it as-is - rather than concatenating per-item tensors into a second full copy. - - That one allocation is the whole of the request's residency, so the byte - budget charges it once, from its first item until the request is - stripped. The scheduler derives the budget from live states each tick, - which makes clearing the state *the* release; there is no release call to - forget or replay. - - The state is rank-local and never crosses a serialization boundary: - schedule distribution carries request/item IDs only, and every rank - constructs its own state at admission. + """Tracks the encoder-cache key for each MM item in a request. + + The state is created when an item-scheduled request enters the executor. + Other requests leave `py_mm_encoder_state` as `None`. + + Encoder outputs live only in the model-owned `TensorLRUCache`. This object + stores one cache key and one ready flag per item, in prompt order. On the + item-scheduled executor, each non-empty slot holds one reference. Two + identical items may point to the same cache entry, but each item still + holds and releases its own reference. The request does not store another + output tensor or combined output buffer. + + Cleanup clears the slots before releasing their references, so repeated + request cleanup is safe. """ embedding_lengths: List[int] - """Declared embedding row count of each atomic item, in prompt order. - `record()` validates incoming tensors against these.""" + """Expected encoder-output rows for each item, in prompt order.""" encoder_token_lengths: List[int] """Validated encoder attention-token cost of each atomic item. @@ -131,23 +130,22 @@ class MultimodalEncoderRequestState: never has to re-validate user input inside the scheduler loop. """ - recorded: List[bool] - """Whether each atomic item has been written into `embeddings`, in prompt - order.""" + item_ready: List[bool] + """Whether each item's bound cache entry is ready, in prompt order.""" + + item_cache_keys: List[Optional[Hashable]] = field(default_factory=list) + """Cache keys held by this request, in prompt order. + + A non-empty slot holds one reference. `item_ready` tells whether the entry + already contains its encoder output. + """ - cache_item_keys: Union[List[Hashable], None, "_Unset"] = _UNSET - """Memoized per-item encoder cache keys, or `None` once known to be - unkeyable. Derived from the request's content hashes and item metadata, - both fixed at admission, so this never goes stale and is computed on the - first iteration that schedules any of the request's items rather than on - every one. `_UNSET` distinguishes "not computed yet" from "computed, and - this request cannot participate in the cache".""" + stable_item_cache_keys: "list[Hashable] | _Unset | None" = _UNSET + """Cached per-item reuse keys. - embeddings: Optional[torch.Tensor] = None - """Contiguous ``[sum(embedding_lengths), ...]`` storage for every item of - this request, allocated by the first `record()`. Published by reference at - `finalize()` and released when the request is stripped, so its - presence is exactly what the byte budget charges.""" + `_UNSET` means the keys have not been computed. `None` means stable keys + cannot be built, so this request uses request-local cache keys instead. + """ @classmethod def from_embedding_lengths( @@ -160,123 +158,72 @@ def from_embedding_lengths( encoder_token_lengths = embedding_lengths return cls(embedding_lengths=list(embedding_lengths), encoder_token_lengths=list(encoder_token_lengths), - recorded=[False] * len(embedding_lengths)) + item_ready=[False] * len(embedding_lengths), + item_cache_keys=[None] * len(embedding_lengths)) def __post_init__(self) -> None: + if not self.item_cache_keys: + self.item_cache_keys = [None] * len(self.item_ready) if not (len(self.embedding_lengths) == len(self.encoder_token_lengths) - == len(self.recorded)): + == len(self.item_ready) == len(self.item_cache_keys)): raise ValueError("MM encoder token and embedding lengths must have " - "exactly one entry per item slot") - # Row offsets are fixed once the declared lengths are known, so derive - # them here rather than re-summing the prefix on every `record()` - # (quadratic in the item count, and every caller now routes through - # `record`). Has one extra entry so `_row_starts[i + 1]` is the end of - # item `i`; the last is the buffer's total row count. - self._row_starts = list( - itertools.accumulate(self.embedding_lengths, initial=0)) + "exactly one cache key per item slot") @property def num_items(self) -> int: - return len(self.recorded) - - @property - def has_storage(self) -> bool: - """Whether this request's embedding buffer is allocated. - - The scheduler reads this to charge the byte budget once per request: - the first scheduled item allocates storage for *all* of them, so - later items of the same request cost nothing more. - """ - return self.embeddings is not None + return len(self.item_ready) @property def progress(self) -> MultimodalEncoderProgress: - if all(self.recorded): + if all(self.item_ready): return MultimodalEncoderProgress.READY - if any(self.recorded): + if any(self.item_ready): return MultimodalEncoderProgress.PARTIAL return MultimodalEncoderProgress.PENDING def pending_item_indices(self) -> List[int]: """Indices of items that still need an encoder output, prompt order.""" return [ - item_idx for item_idx, done in enumerate(self.recorded) if not done + item_idx for item_idx, ready in enumerate(self.item_ready) + if not ready ] - def record(self, item_idx: int, output: torch.Tensor) -> None: - """Copy one item's encoder output into its row range of the buffer. - - The first call allocates the buffer for every item of the request, - which is why it is deferred to here: the row shape, dtype and device - only become known with the first encoder output. ``output`` may be a - view of a batched encoder output or a read-through cache entry, and - the copy is the single choke point that keeps the request from - retaining the whole batch allocation through a view or aliasing a - cache entry that can be evicted under it. - - Raises when the tensor does not match the item's declared row count, - when the item was already recorded, or when it disagrees with the - buffer on trailing shape/dtype/device (items of one request share one - contiguous embedding). - """ - expected_rows = self.embedding_lengths[item_idx] - if output.shape[0] != expected_rows: - raise MultimodalEncoderRequestError( - f"MM item {item_idx} produced {output.shape[0]} embeddings; " - f"expected {expected_rows}") - if self.recorded[item_idx]: - raise MultimodalEncoderRequestError( - f"MM item {item_idx} was already recorded; items are " - "encoded at most once per request") - if self.embeddings is None: - self.embeddings = torch.empty( - (self._row_starts[-1], *output.shape[1:]), - dtype=output.dtype, - device=output.device) - elif (self.embeddings.shape[1:] != output.shape[1:] - or self.embeddings.dtype != output.dtype - or self.embeddings.device != output.device): - raise MultimodalEncoderRequestError( - "MM encoder items for one request must have matching " - "output shape, dtype, and device") - start = self._row_starts[item_idx] - self.embeddings[start:start + expected_rows].copy_(output.detach()) - self.recorded[item_idx] = True - - def resident_output_bytes(self, bytes_per_encoder_embedding: int) -> int: - """Bytes of encoder output this request holds on the device. - - The scheduler sums this over live states every tick to derive the - occupied share of the encoder output byte budget; there is no - separate accounting to keep in sync. - - The buffer covers every item from the first `record()` onward, so a - partially encoded request already charges its full footprint — which - is what it actually occupies — and keeps charging it after - `finalize()` until the request is stripped post-prefill. - """ - if self.embeddings is None: - return 0 - return self._row_starts[-1] * bytes_per_encoder_embedding - - def finalize(self, multimodal_data: Dict[str, Any]) -> bool: - """Publish the request's embedding once every item is recorded. - - Attaches the buffer as ``multimodal_embedding`` and drops the raw - pre-encoder inputs. The buffer is already the contiguous form the - prefill path wants, so publishing is a reference — nothing is copied - here and nothing downstream concatenates the items back together. - - The state keeps its own reference so `resident_output_bytes()` still - charges the rows, which stay resident until the request is stripped - post-prefill. Both references die together at that strip. - No-op returning ``False`` while any item is still pending. - """ - if not self.recorded or not all(self.recorded): - return False - multimodal_data["multimodal_embedding"] = self.embeddings - strip_mm_encoder_inputs(multimodal_data) - return True + def set_item_cache_key(self, item_idx: int, cache_key: Hashable, *, + ready: bool) -> None: + """Assign a cache key to one item.""" + if self.item_cache_keys[item_idx] is not None: + raise RuntimeError(f"MM item {item_idx} already has a cache key") + self.item_cache_keys[item_idx] = cache_key + self.item_ready[item_idx] = ready + + def clear_item_cache_key(self, item_idx: int) -> Hashable: + """Clear and return one item's cache key.""" + cache_key = self.item_cache_keys[item_idx] + if cache_key is None: + raise RuntimeError(f"MM item {item_idx} has no cache key") + self.item_cache_keys[item_idx] = None + self.item_ready[item_idx] = False + return cache_key + + def mark_cache_key_ready(self, cache_key: Hashable) -> None: + """Mark every item using `cache_key` as ready.""" + for item_idx, bound_cache_key in enumerate(self.item_cache_keys): + if bound_cache_key == cache_key and not self.item_ready[item_idx]: + self.item_ready[item_idx] = True + + def mark_all_items_ready(self) -> None: + """Mark every bound item ready after a scheduled context replay.""" + self.item_ready = [True] * self.num_items + + def pop_all_cache_keys(self) -> List[Hashable]: + """Return and clear all cache keys, keeping prompt order and duplicates.""" + cache_keys = [ + cache_key for cache_key in self.item_cache_keys + if cache_key is not None + ] + self.item_cache_keys = [None] * len(self.item_cache_keys) + self.item_ready = [False] * len(self.item_ready) + return cache_keys if TYPE_CHECKING: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index f23325af2101..8d04d101f880 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -11,8 +11,8 @@ import weakref from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import (Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, - Union) +from typing import (Any, Callable, Dict, Hashable, List, Optional, Sequence, + Tuple, Type, Union) import torch import torch._dynamo.config @@ -76,6 +76,7 @@ from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from ..speculative.interface import INVALID_PROMPT_LOOKAHEAD_TOKEN from ..speculative.spec_sampler_base import SampleStateTensorsSpec +from ..tensor_lru_cache import TensorLRUCache from ..utils import (get_model_extra_attrs, get_per_request_prefill_cuda_graph_flag, set_per_request_prefill_cuda_graph_flag, @@ -557,6 +558,17 @@ def __init__( # Absent, not None, when item scheduling is off: external readers # rely on the `getattr` default. self.bytes_per_mm_encoder_embedding = mm_item_scheduler.bytes_per_embedding + if isinstance(self.model, + MultimodalModelMixin) and mapping.is_first_pp_rank(): + multimodal_config = self.model.model_config.multimodal_config + reuse_capacity_bytes = (multimodal_config.encoder_cache_max_bytes + if self.model.encoder_cache_active else 0) + cache_capacity_bytes = max( + self.mm_encoder_output_budget_bytes or 0, + reuse_capacity_bytes, + ) + self.model._initialize_multimodal_encoder_cache( + cache_capacity_bytes) setup_mm_encoder_attn_metadata( self.model, self.input_processor, self.encoder_max_num_tokens, mm_item_scheduler.attention_metadata_capacity @@ -3563,7 +3575,7 @@ def forward_multimodal_encoder_items( requests: List[LlmRequest], scheduled_items: Dict[int, List[int]], ) -> None: - """Forward selected MM encoder items and commit request-local outputs.""" + """Forward selected MM encoder items into the unified output cache.""" if not scheduled_items: return if self._mm_item_scheduler is None: @@ -3571,6 +3583,35 @@ def forward_multimodal_encoder_items( "Item-level MM scheduling requires MultimodalModelMixin") self._mm_item_scheduler.forward_items(requests, scheduled_items) + @property + def mm_encoder_cache(self) -> Optional[TensorLRUCache]: + """Return the model-owned encoder-output cache used by this engine.""" + if not isinstance(self.model, MultimodalModelMixin): + return None + if self._mm_item_scheduler is None and not self.model.encoder_cache_active: + return None + return self.model._multimodal_encoder_cache + + def get_mm_encoder_item_cache_keys( + self, request: LlmRequest) -> Optional[List[Hashable]]: + """Return stable cache keys for an item-scheduled request, when available.""" + if self._mm_item_scheduler is None: + return None + return self._mm_item_scheduler.item_cache_keys(request) + + def invalidate_multimodal_encoder_cache(self) -> None: + """Clear cached MM encoder outputs when no request is using them.""" + encoder_cache = self.mm_encoder_cache + if encoder_cache is not None: + encoder_cache.clear() + + def _build_multimodal_data_for_llm( + self, request: LlmRequest) -> Optional[Dict[str, Any]]: + """Attach cached item outputs when item scheduling owns the request.""" + if self._mm_item_scheduler is None: + return request.py_multimodal_data + return self._mm_item_scheduler.build_multimodal_data_for_llm(request) + def _set_up_spec_metadata( self, spec_resource_manager: Optional[BaseResourceManager], @@ -5484,7 +5525,7 @@ def append_cross_attention_state(request: LlmRequest, multimodal_params = MultimodalParams( multimodal_input=_build_request_multimodal_input( request, self._mm_encoder_cache_enabled), - multimodal_data=request.py_multimodal_data, + multimodal_data=self._build_multimodal_data_for_llm(request), multimodal_runtime=py_multimodal_runtime, mm_item_order=getattr(request, "py_mm_item_order", None), input_ids_start_offset=context_start_idx) @@ -6627,7 +6668,8 @@ def _prepare_tp_inputs_no_cache( multimodal_params = MultimodalParams( multimodal_input=_build_request_multimodal_input( request, self._mm_encoder_cache_enabled), - multimodal_data=request.py_multimodal_data, + multimodal_data=self._build_multimodal_data_for_llm( + request), mm_item_order=getattr(request, "py_mm_item_order", None), input_ids_start_offset=context_start_idx) multimodal_params.to_device("multimodal_data", diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 339f7ecb26bc..f9c40d6458a0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -293,26 +293,6 @@ def _load_iteration_indexes(env_var: str): return frozenset(starts), frozenset(stops) -def _strip_py_multimodal_data_post_prefill(request: LlmRequest) -> None: - """Drop encoder outputs and raw inputs after prefill or early termination. - - Wraps `strip_mm_data_for_generation` and mutates the shared `request.py_multimodal_data` - in-place so the `LlmRequest`'s multimodal tensors actually get freed (unlike - `MultimodalParams.strip_for_generation`, which rebinds a per-forward-call wrapper's attribute - and leaves the request's dict untouched). - """ - mm_data = getattr(request, "py_multimodal_data", None) - if mm_data: - strip_mm_data_for_generation(mm_data) - # Drop the per-item encoder state alongside the dict clear above. The - # state and the published `multimodal_embedding` are two references to - # one embedding buffer, so both have to go for the GPU memory to be - # freed -- and a request stripped mid-encode only has the state's. - # Clearing it is also the byte-budget release: the scheduler derives - # occupancy from live states, so this request stops counting next tick. - request.py_mm_encoder_state = None - - @contextmanager def _distributed_warmup_guard(dist: Distributed, mapping: Mapping) -> Iterator[None]: @@ -6518,10 +6498,8 @@ def _schedule(self): def _forward_multimodal_encoder_step( self, scheduled_requests: ScheduledRequests) -> None: - """Run scheduler-selected MM encoder work before LLM resources.""" - scheduled_items = scheduled_requests.scheduled_mm_encoder_items - if not scheduled_items: - return + """Update the MM cache and encode selected items before LLM prefill.""" + scheduled_items = scheduled_requests.scheduled_mm_encoder_items or {} try: self.model_engine.forward_multimodal_encoder_items( self.active_requests, scheduled_items) @@ -6530,25 +6508,32 @@ def _forward_multimodal_encoder_step( logger.error(f"Encountered an error in multimodal encoder forward: " f"{error_msg}\n{traceback.format_exc()}") - failed_request_ids = set(scheduled_items) - failed_requests = [ - request for request in self.active_requests - if request.request_id in failed_request_ids - ] + failed_request_ids = set(e.request_ids or scheduled_items) + self._handle_multimodal_encoder_request_error( + scheduled_requests, error_msg, failed_request_ids) + + def _handle_multimodal_encoder_request_error( + self, scheduled_requests: ScheduledRequests, error_msg: str, + failed_request_ids: set[int]) -> None: + """Remove requests that depend on a failed MM encoder output.""" + failed_requests = [ + request for request in self.active_requests + if request.request_id in failed_request_ids + ] - # Capacity scheduling may already have placed requests whose last - # pending item was selected into this iteration's context batch. - # Remove the failed owners before the LLM forward; unrelated - # context and generation work remains intact. - scheduled_requests.reset_context_requests([ - request for request in scheduled_requests.context_requests - if request.request_id not in failed_request_ids - ]) - scheduled_requests.scheduled_mm_encoder_items = None + # Capacity scheduling may already have placed requests whose last + # pending item was selected into this iteration's context batch. + # Remove the failed owners before the LLM forward; unrelated context + # and generation work remains intact. + scheduled_requests.reset_context_requests([ + request for request in scheduled_requests.context_requests + if request.request_id not in failed_request_ids + ]) + scheduled_requests.scheduled_mm_encoder_items = None - self._handle_errors(error_msg, - requests=failed_requests, - charge_budget=False) + self._handle_errors(error_msg, + requests=failed_requests, + charge_budget=False) # --------------------------------------------------------------- # Encoder-decoder support: encoder iteration in the executor loop. @@ -8030,6 +8015,22 @@ def _update_generation_requests_that_will_complete_next_iteration( request.set_exclude_last_generation_logits(False) request.state = LlmRequestState.GENERATION_TO_COMPLETE + def _release_multimodal_resources(self, request: LlmRequest) -> None: + """Release this request's cache entries and discard unused MM data.""" + state = request.py_mm_encoder_state + if state is not None: + cache_keys = state.pop_all_cache_keys() + request.py_mm_encoder_state = None + encoder_cache = self.model_engine.mm_encoder_cache + if encoder_cache is None: + raise RuntimeError("MM request state requires an encoder cache") + for cache_key in cache_keys: + encoder_cache.release(cache_key) + + mm_data = request.py_multimodal_data + if mm_data: + strip_mm_data_for_generation(mm_data) + def _update_request_states_tp(self, scheduled_requests: ScheduledRequests): # handle potential attention dp dummy request if self.active_requests and self.active_requests[ @@ -8053,7 +8054,7 @@ def _update_request_states_tp(self, scheduled_requests: ScheduledRequests): # on `py_multimodal_data`. Without this, encoder inputs and outputs for multi-modal # requests stay pinned on GPU through the full decode lifetime and can lead to OOMs # at high concurrency. - _strip_py_multimodal_data_post_prefill(request) + self._release_multimodal_resources(request) if not self.disable_overlap_scheduler and request.will_complete_next_iteration( ): request.set_exclude_last_generation_logits(False) @@ -8384,7 +8385,7 @@ def _do_terminate_request(self, request: LlmRequest) -> None: self._free_request_resources(request) # Cancellation and request-scoped failures can terminate before the # normal post-prefill release point, including with a partial buffer. - _strip_py_multimodal_data_post_prefill(request) + self._release_multimodal_resources(request) if self.gather_all_responses or self.dist.rank == 0: self.result_wait_queues.pop(request.py_request_id, None) @@ -8927,6 +8928,20 @@ def _handle_speculative_decoding( def reset_prefix_cache(self): self.kv_cache_manager.reset_reuse_state() + def invalidate_multimodal_encoder_cache(self) -> None: + """Clear weight-derived MM outputs after request lifetimes quiesce.""" + live_request_ids = [ + request.py_request_id for request in self.active_requests + if request.py_mm_encoder_state is not None and any( + cache_key is not None + for cache_key in request.py_mm_encoder_state.item_cache_keys) + ] + if live_request_ids: + raise RuntimeError( + "cannot update weights with live multimodal cache references: " + f"request_ids={live_request_ids}") + self.model_engine.invalidate_multimodal_encoder_cache() + def _handle_guided_decoder_errors( self, scheduled_batch: ScheduledRequests, failed_requests: Optional[List[Tuple[int, str]]]): diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 6f3b67f9459d..76511cb7874c 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -7,9 +7,10 @@ import inspect from abc import ABC, abstractmethod from collections import namedtuple +from collections.abc import Hashable from dataclasses import dataclass from enum import Enum -from typing import Any, Callable, Optional, TypeAlias, TypeVar +from typing import Any, Callable, Optional, TypeAlias, TypeVar, cast from strenum import StrEnum @@ -17,12 +18,15 @@ from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy from tensorrt_llm.logger import logger +from ...tensor_lru_cache import CacheAcquireResult, TensorLRUCache + # Assuming these imports exist in your environment from ..llm_request import ( LlmRequest, LlmRequestState, - format_multimodal_encoder_output_budget_error, + MultimodalEncoderRequestState, is_multimodal_encoder_ready, + make_mm_encoder_transient_cache_key, ) RequestList = list[LlmRequest] @@ -77,6 +81,7 @@ class SchedulerOutput( "paused_requests", "fitting_disagg_gen_init_requests", "num_fitting_requests", + # Multimodal encoder scheduling outputs. "scheduled_mm_encoder_items", "recompute_paused_requests", ], @@ -84,7 +89,7 @@ class SchedulerOutput( ): """Scheduler result. - ``scheduled_mm_encoder_items`` defaults to ``None``. The V2-only + The optional multimodal fields default to ``None``. The V2-only ``recompute_paused_requests`` defaults to a fresh empty list so existing V1 schedulers can keep constructing the original six-field output. """ @@ -570,7 +575,7 @@ def can_schedule(self, requests: RequestList) -> bool: class MultimodalScheduler(RequestScheduler): - """Add atomic multimodal item budgeting around the existing scheduler. + """Add per-item MM encoder limits and cache lookup to the LLM scheduler. The wrapper is constructed only for ``MultimodalModelMixin`` models. It deliberately reuses the wrapped scheduler's capacity and microbatch @@ -582,12 +587,11 @@ class MultimodalScheduler(RequestScheduler): attention sequences. Those attention metadata capacities are derived separately from the token budget and model geometry. - When ``output_budget_bytes`` is configured, selection also enforces the - encoder output byte budget (allocate-before-compute): an item is only - selected when its embedding bytes fit alongside the outputs already - resident on live requests and bytes claimed earlier in the pass. - Occupancy is derived from request states each pass rather than tracked - by a counter, so a stripped or aborted request self-heals the budget. + Before using encoder budget, it acquires cache entries in prompt order. + Cache hits and items already being encoded need no additional encoder + work. Each new reservation selects one item to encode. The cache tracks + bytes and references, retains reusable outputs, and chooses LRU entries to + remove when space is needed. """ def __init__( @@ -596,26 +600,18 @@ def __init__( max_batch_size: int, max_num_tokens: int, *, - output_budget_bytes: int | None = None, - bytes_per_encoder_embedding: int = 0, + encoder_cache: TensorLRUCache[Hashable], + get_item_cache_keys: Callable[[LlmRequest], list[Hashable] | None], + bytes_per_encoder_embedding: int, + retain_cache_entries: bool, ) -> None: self.scheduler = scheduler self.max_batch_size = max_batch_size self.max_num_tokens = max_num_tokens - # Optional byte budget for encoder outputs living outside a forward - # pass. Item selection performs allocate-before-compute against it: - # occupancy is *derived* each pass from live request states (their - # recorded, not-yet-consumed outputs) — there is no counter to - # release or keep in sync; a stripped or aborted request simply - # stops contributing. `bytes_per_encoder_embedding` converts declared - # embedding rows to bytes and must be positive alongside a budget. - self.output_budget_bytes = output_budget_bytes + self.encoder_cache = encoder_cache + self.get_item_cache_keys = get_item_cache_keys self.bytes_per_encoder_embedding = bytes_per_encoder_embedding - if output_budget_bytes is not None and bytes_per_encoder_embedding <= 0: - raise ValueError( - "bytes_per_encoder_embedding must be positive when a byte " - "budget bounds MM encoder outputs" - ) + self.retain_cache_entries = retain_cache_entries self.has_separate_stages = hasattr(scheduler, "capacity_scheduler") and hasattr( scheduler, "micro_batch_scheduler" ) @@ -624,23 +620,55 @@ def __init__( def scheduling_state_range(self) -> tuple[LlmRequestState, LlmRequestState]: return self.scheduler.scheduling_state_range - def _total_resident_output_bytes(self, active_requests: RequestList) -> int: - """Sum per-request resident encoder-output bytes across live states. + def _acquire_request_cache_entries( + self, request: LlmRequest, state: MultimodalEncoderRequestState + ) -> list[int] | None: + """Acquire every missing item cache entry, or undo the attempt if full.""" + stable_keys = self.get_item_cache_keys(request) + if stable_keys is None: + # Give items without reusable keys a temporary key so they can use + # the same cache path. + item_cache_keys = [ + make_mm_encoder_transient_cache_key(request.request_id, item_idx) + for item_idx in range(state.num_items) + ] + retain_after_release = False + else: + item_cache_keys = stable_keys + retain_after_release = self.retain_cache_entries + acquired_item_indices: list[int] = [] + for item_idx in state.pending_item_indices(): + if state.item_cache_keys[item_idx] is not None: + continue + cache_key = item_cache_keys[item_idx] + expected_bytes = state.embedding_lengths[item_idx] * self.bytes_per_encoder_embedding + acquire_result = self.encoder_cache.acquire( + cache_key, + expected_bytes, + retain_after_release=retain_after_release, + ) + if acquire_result is None: + self._release_acquired_cache_entries(state, acquired_item_indices) + return None + state.set_item_cache_key( + item_idx, + cache_key, + ready=acquire_result is CacheAcquireResult.READY_HIT, + ) + acquired_item_indices.append(item_idx) + return acquired_item_indices - Summed fresh every pass: a request whose outputs were consumed - (stripped post-prefill) or that was aborted no longer contributes, - so the accounting self-heals with no release bookkeeping. - """ - return sum( - state.resident_output_bytes(self.bytes_per_encoder_embedding) - for request in active_requests - if (state := request.py_mm_encoder_state) is not None - ) + def _release_acquired_cache_entries( + self, + state: MultimodalEncoderRequestState, + acquired_item_indices: list[int], + ) -> None: + for item_idx in acquired_item_indices: + cache_key = state.clear_item_cache_key(item_idx) + self.encoder_cache.release(cache_key) - def _select_items( - self, requests: RequestList, *, active_requests: RequestList | None = None - ) -> tuple[dict[int, list[int]], RequestList]: - """Greedily select pending MM items under the encoder budgets. + def _select_items(self, requests: RequestList) -> tuple[dict[int, list[int]], RequestList]: + """Acquire output entries and select the items that need encoding. Requests are visited in the wrapped capacity scheduler's FCFS order with no explicit `MultimodalEncoderProgress`-based priority: a @@ -648,96 +676,74 @@ def _select_items( anything admitted later, so its remaining items resume before newer work by order alone. - When a byte budget is configured, selection also performs - allocate-before-compute, per request rather than per item: a request - starts only if its *whole* embedding fits alongside (a) storage - already held by live requests (derived from `active_requests`) and - (b) bytes claimed earlier in this pass. That matches how the storage - is allocated — the first recorded item sizes the buffer for all of - them — and means a started request can always finish, so no - head-of-line reservation is needed to keep later requests from - squatting the space it still needs. - - Returns the selected item indices per request id, plus the requests - eligible for LLM microbatch scheduling this iteration (encoder - outputs already ready, or every pending item selected above). + Cache hits and items already selected through the same cache entry use + no encoder compute budget. Newly acquired entries are kept when the + request becomes ready from cache hits or a pending output will be + produced this iteration. After producer selection, their output bytes + are summed and passed once to `ensure_capacity` before producer + commits. """ remaining_batch_slots = self.max_batch_size remaining_tokens = self.max_num_tokens - budget = self.output_budget_bytes - resident_bytes = ( - self._total_resident_output_bytes( - active_requests if active_requests is not None else requests - ) - if budget is not None - else 0 - ) - reserved_bytes = 0 selected: dict[int, list[int]] = {} - llm_eligible: RequestList = [] + selected_cache_keys: set[Hashable] = set() + selected_output_bytes = 0 for request in requests: state = request.py_mm_encoder_state - if state is None: - llm_eligible.append(request) + if state is None or is_multimodal_encoder_ready(request): continue - if is_multimodal_encoder_ready(request): - llm_eligible.append(request) + acquired_item_indices = self._acquire_request_cache_entries(request, state) + if acquired_item_indices is None: continue - - # Admission validates user-provided item metadata and stores an - # owned copy on the request state. Reading only that state here - # keeps malformed-input failures scoped to the affected request - # instead of raising from the scheduler loop. - token_lengths = state.encoder_token_lengths - - pending = state.pending_item_indices() - # The first item scheduled for a request allocates the storage for - # *all* of its items, so the byte budget is charged once per - # request rather than per item. A request that cannot be charged - # yet stays fully pending instead of occupying part of the budget - # with work that cannot be prefilled until it completes. - if ( - budget is not None - and not state.has_storage - and pending - and remaining_batch_slots > 0 - and token_lengths[pending[0]] <= remaining_tokens - ): - request_bytes = sum(state.embedding_lengths) * self.bytes_per_encoder_embedding - if request_bytes > budget: - # Liveness backstop: admission - # (`initialize_multimodal_encoder_request`) already - # rejects requests whose outputs can never coexist - # within the budget, so reaching this means an - # accounting bug rather than a user input. - raise RuntimeError( - format_multimodal_encoder_output_budget_error( - request_bytes, - budget, - self.max_num_tokens, - request_id=request.py_request_id, - ) - ) - if resident_bytes + reserved_bytes + request_bytes > budget: - continue - reserved_bytes += request_bytes - + pending_items = [ + (item_idx, cast(Hashable, state.item_cache_keys[item_idx])) + for item_idx in state.pending_item_indices() + ] + will_make_progress = any( + cache_key in selected_cache_keys for _, cache_key in pending_items + ) request_items: list[int] = [] - for item_idx in pending: - cost = token_lengths[item_idx] + for item_idx, cache_key in pending_items: + if cache_key in selected_cache_keys: + continue + cost = state.encoder_token_lengths[item_idx] if remaining_batch_slots == 0 or cost > remaining_tokens: break request_items.append(item_idx) + selected_cache_keys.add(cache_key) + will_make_progress = True + selected_output_bytes += ( + state.embedding_lengths[item_idx] * self.bytes_per_encoder_embedding + ) remaining_batch_slots -= 1 remaining_tokens -= cost if request_items: selected[request.request_id] = request_items - if pending and len(request_items) == len(pending): + # Release newly acquired entries if none of this request's pending + # outputs will be produced in this iteration. + if acquired_item_indices and pending_items and not will_make_progress: + self._release_acquired_cache_entries(state, acquired_item_indices) + + llm_eligible: RequestList = [] + # Check after all item selection because one request may select a + # shared cache entry needed by another request. + for request in requests: + state = request.py_mm_encoder_state + if state is None: + llm_eligible.append(request) + continue + ready_after_encoder_step = all( + ready or (cache_key is not None and cache_key in selected_cache_keys) + for ready, cache_key in zip(state.item_ready, state.item_cache_keys, strict=True) + ) + if ready_after_encoder_step: llm_eligible.append(request) + self.encoder_cache.ensure_capacity(selected_output_bytes) + return selected, llm_eligible def _schedule_micro_batch( @@ -785,8 +791,7 @@ def schedule_request( active_requests, inflight_request_ids ) selected_items, llm_eligible = self._select_items( - list(scheduler_output.context_requests), - active_requests=active_requests, + list(scheduler_output.context_requests) ) return scheduler_output._replace( context_requests=llm_eligible, @@ -798,9 +803,7 @@ def schedule_request( fitting_requests, fitting_disagg_gen_init_requests, paused_requests = ( self.scheduler.capacity_scheduler.schedule_request(active_requests) ) - selected_items, llm_eligible = self._select_items( - list(fitting_requests), active_requests=active_requests - ) + selected_items, llm_eligible = self._select_items(list(fitting_requests)) # Preserve the capacity scheduler's decisions while attaching the MM # item plan that the executor must run before the selected LLM # microbatch. @@ -833,7 +836,9 @@ def schedule_request( if not self.has_separate_stages: scheduler_output = self.scheduler.schedule_request(llm_eligible, inflight_request_ids) - return scheduler_output._replace(scheduled_mm_encoder_items=selected_items or None) + return scheduler_output._replace( + scheduled_mm_encoder_items=selected_items or None, + ) llm_eligible_ids = {request.request_id for request in llm_eligible} fitting_requests, fitting_disagg_gen_init_requests, paused_requests = ( diff --git a/tensorrt_llm/_torch/tensor_lru_cache.py b/tensorrt_llm/_torch/tensor_lru_cache.py index 58e789e86309..ae5e8353ad9e 100644 --- a/tensorrt_llm/_torch/tensor_lru_cache.py +++ b/tensorrt_llm/_torch/tensor_lru_cache.py @@ -19,6 +19,7 @@ from collections import OrderedDict from collections.abc import Hashable from dataclasses import dataclass +from enum import Enum, auto from threading import RLock from typing import Generic, NamedTuple, TypeVar @@ -29,18 +30,44 @@ K = TypeVar("K", bound=Hashable) -class _Entry(NamedTuple): - value: torch.Tensor +class _CacheEntryState(Enum): + """Lifecycle state of a cache key.""" + + # Expected bytes and references exist, but no tensor has been committed. + RESERVED = auto() + # The cache owns a tensor that callers may read. + READY = auto() + + +class CacheAcquireResult(Enum): + """How `acquire` satisfied a successful request.""" + + # A READY entry was found and referenced. + READY_HIT = auto() + # A missing key became a new RESERVED entry. + NEW_RESERVATION = auto() + # An existing RESERVED entry gained another reference. + RESERVATION_HIT = auto() + + +@dataclass +class _Entry: + state: _CacheEntryState size_bytes: int + reference_count: int + retain_after_release: bool + value: torch.Tensor | None = None # CUDA event recorded on the producing stream right after the clone in `put`. Consumers on a # different stream wait on it before reading `value`. `None` for CPU tensors or when the cache # is not stream-aware. - producer_event: torch.cuda.Event | None + producer_event: torch.cuda.Event | None = None class TensorLRUCacheStats(NamedTuple): max_bytes: int current_bytes: int + reserved_bytes: int + in_use_bytes: int item_count: int hits: int misses: int @@ -48,6 +75,8 @@ class TensorLRUCacheStats(NamedTuple): replacements: int evictions: int rejected_insertions: int + producer_misses: int + inflight_deduplications: int hit_rate: float @@ -59,6 +88,8 @@ class _CacheCounters: replacements: int = 0 evictions: int = 0 rejected_insertions: int = 0 + producer_misses: int = 0 + inflight_deduplications: int = 0 @property def hit_rate(self) -> float: @@ -81,6 +112,24 @@ class TensorLRUCache(Generic[K]): temporarily needs both the source tensor and its copy and may exceed the cache limit until eviction completes. + Managed entries add a RESERVED state and reference-counted READY state to the original cache: + + 1. `acquire` adds a reference and creates a `RESERVED` entry on a miss. + 2. `ensure_capacity` evicts only unreferenced `READY` entries before selected + producers materialize their outputs. + 3. `commit` stores a producer output in its `RESERVED` entry without + choosing more eviction victims. + 4. `get` reads a `READY` tensor without changing its reference count. + 5. `release` drops the reference and applies the entry's retention policy. + + Reservations and referenced READY entries cannot be evicted. An entry is + referenced after `acquire` increments its `reference_count` and until the + matching `release` calls reduce it to zero. `in_use_bytes` is the logical + size of those referenced READY tensors. + + `pop` is intentionally separate from `release`: it physically removes an + unreferenced entry without applying reference or retention policy. + In CUDA-stream-aware mode, each entry owns the event recorded after its clone. Replacement, eviction, and clear drop that event with the entry; events are not reused because an evicted tensor may still have outstanding consumers on another stream. @@ -107,6 +156,9 @@ def __init__( self._name = name self._cuda_stream_aware = cuda_stream_aware self._current_bytes = 0 + self._reserved_bytes = 0 + # READY tensor bytes whose reference_count is positive, so they cannot be evicted. + self._in_use_bytes = 0 self._items: OrderedDict[K, _Entry] = OrderedDict() self._lock = RLock() self._counters = _CacheCounters() @@ -124,27 +176,116 @@ def __len__(self) -> int: with self._lock: return len(self._items) - def get(self, key: K) -> torch.Tensor | None: + def acquire( + self, + key: K, + expected_bytes: int, + *, + retain_after_release: bool = True, + ) -> CacheAcquireResult | None: + """Acquire a reference, reserving a missing entry when needed. + + A READY hit and an existing reservation both gain one reference. A + missing key becomes RESERVED and accounts `expected_bytes` until the + producer stores its tensor with `commit`. + + Args: + key: Identity shared by producers and consumers of one tensor. + expected_bytes: Exact tensor bytes that a missing key will produce. + retain_after_release: Whether a READY entry remains reusable after + its final reference is released. Stable content-key entries + use `True` for cross-request reuse; request-local entries use + `False` and are removed immediately. + + Returns: + How the reference was acquired, or `None` when live references and + reservations temporarily consume all capacity. + + Raises: + ValueError: If the requested size is invalid or existing metadata + conflicts with the request. + """ + if expected_bytes <= 0: + raise ValueError("expected_bytes must be positive") + if expected_bytes > self._max_bytes: + raise ValueError( + f"expected_bytes ({expected_bytes}) exceeds cache capacity ({self._max_bytes})" + ) + + with self._lock: + entry = self._items.get(key) + if entry is not None: + if entry.size_bytes != expected_bytes: + raise ValueError( + f"existing cache entry size ({entry.size_bytes}) does not match " + f"expected_bytes ({expected_bytes})" + ) + if entry.retain_after_release != retain_after_release: + raise ValueError("existing cache entry retention policy does not match") + + if entry.state is _CacheEntryState.RESERVED: + entry.reference_count += 1 + self._counters.inflight_deduplications += 1 + return CacheAcquireResult.RESERVATION_HIT + if entry.state is not _CacheEntryState.READY: + raise RuntimeError(f"unexpected cache entry state: {entry.state}") + + if entry.reference_count == 0: + # A retained cache hit becomes non-evictable again. Make + # sure it fits beside current reservations and references + # before claiming it for this request. + if self._claimed_bytes + entry.size_bytes > self._max_bytes: + return None + self._in_use_bytes += entry.size_bytes + entry.reference_count += 1 + self._items.move_to_end(key) + self._counters.hits += 1 + return CacheAcquireResult.READY_HIT + + if self._claimed_bytes + expected_bytes > self._max_bytes: + return None + + self._items[key] = _Entry( + state=_CacheEntryState.RESERVED, + size_bytes=expected_bytes, + reference_count=1, + retain_after_release=retain_after_release, + ) + self._reserved_bytes += expected_bytes + self._counters.misses += 1 + self._counters.producer_misses += 1 + return CacheAcquireResult.NEW_RESERVATION + + def get(self, key: K, *, record_stats: bool = True) -> torch.Tensor | None: """Return a cache-owned, immutable tensor and promote it to most-recently-used. The returned tensor aliases the cached value. Callers must not mutate it. + Only READY entries are returned. `get` does not acquire a managed + reference; `record_stats=False` suppresses its hit/miss accounting. """ with self._lock: entry = self._items.get(key) - if entry is None: - self._counters.misses += 1 + if entry is None or entry.state is not _CacheEntryState.READY: + if record_stats: + self._counters.misses += 1 return None - self._counters.hits += 1 + if record_stats: + self._counters.hits += 1 self._items.move_to_end(key) self._prepare_for_current_stream(entry) + assert entry.value is not None return entry.value - def put(self, key: K, value: torch.Tensor) -> bool: + def put( + self, + key: K, + value: torch.Tensor, + ) -> bool: """Insert or replace a tensor. - Returns `False` and leaves the cache unchanged when `value` is larger than the full - cache capacity. + Returns `False` and leaves the cache contents unchanged when the value + cannot fit beside current reservations and references. """ size_bytes = self._tensor_size_bytes(value) @@ -157,50 +298,170 @@ def put(self, key: K, value: torch.Tensor) -> bool: ) return False - stored_value = value.detach().clone() - producer_event = None - if self._cuda_stream_aware and stored_value.is_cuda: - producer_event = torch.cuda.Event() - producer_event.record(torch.cuda.current_stream(stored_value.device)) + stored_value, producer_event = self._clone_for_storage(value) with self._lock: - old_entry = self._items.pop(key, None) + old_entry = self._items.get(key) if old_entry is not None: + if old_entry.state is not _CacheEntryState.READY: + raise RuntimeError("cannot replace a reserved cache entry") + if old_entry.reference_count: + raise RuntimeError("cannot replace a referenced cache entry") + + # `put` creates an unreferenced READY entry, but it must leave + # enough room for every reservation and referenced entry to become + # resident. Reject before changing the cache when those protected + # bytes already claim the remaining capacity. + if self._claimed_bytes + size_bytes > self._max_bytes: + self._counters.rejected_insertions += 1 + return False + + if old_entry is not None: + del self._items[key] self._current_bytes -= old_entry.size_bytes self._counters.replacements += 1 else: self._counters.insertions += 1 + # The protected-byte check above guarantees that removable READY + # entries can make this space without failure. + self.ensure_capacity(size_bytes) + self._items[key] = _Entry( - value=stored_value, + state=_CacheEntryState.READY, size_bytes=size_bytes, + reference_count=0, + retain_after_release=True, + value=stored_value, producer_event=producer_event, ) self._current_bytes += size_bytes + return True - evicted_count, evicted_bytes = self._evict_until_within_limit() - if evicted_count: - self._counters.evictions += evicted_count - logger.debug( - f"{self._name}: evicted {evicted_count} LRU entries, " - f"freed_bytes={evicted_bytes}, current_bytes={self._current_bytes}, " - f"max_bytes={self._max_bytes}" + def commit(self, key: K, value: torch.Tensor) -> None: + """Store a producer output in an acquired reservation. + + The cache owns the reservation size and verifies the producer output + before changing the entry to READY. Capacity must already have been + prepared with `ensure_capacity`. + """ + size_bytes = self._tensor_size_bytes(value) + with self._lock: + entry = self._items.get(key) + if entry is None or entry.state is not _CacheEntryState.RESERVED: + raise RuntimeError("cache commit requires a reserved entry") + if size_bytes != entry.size_bytes: + raise ValueError( + f"tensor size ({size_bytes}) does not match reserved bytes ({entry.size_bytes})" ) - return True + if self._current_bytes + size_bytes > self._max_bytes: + raise RuntimeError("reserved entry does not have enough cache space") + + stored_value, producer_event = self._clone_for_storage(value) + entry.state = _CacheEntryState.READY + entry.value = stored_value + entry.producer_event = producer_event + self._reserved_bytes -= size_bytes + self._in_use_bytes += size_bytes + self._current_bytes += size_bytes + self._items.move_to_end(key) + self._counters.insertions += 1 + + def ensure_capacity(self, incoming_bytes: int) -> None: + """Ensure physical space for selected outputs before they are produced. + + `incoming_bytes` is the total size of outputs selected for the next + producer commits. Only unreferenced READY entries are eligible LRU victims; + reservations and referenced tensors are never removed. A failure + leaves the cache unchanged. + """ + if incoming_bytes < 0: + raise ValueError("incoming_bytes must be non-negative") + with self._lock: + required_bytes = max( + 0, + self._current_bytes + incoming_bytes - self._max_bytes, + ) + if required_bytes == 0: + return + + entries_to_evict: list[tuple[K, _Entry]] = [] + freed_bytes = 0 + for cache_key, entry in list(self._items.items()): + if entry.state is not _CacheEntryState.READY or entry.reference_count != 0: + continue + entries_to_evict.append((cache_key, entry)) + freed_bytes += entry.size_bytes + if freed_bytes >= required_bytes: + break + + if freed_bytes < required_bytes: + raise RuntimeError("cache does not have enough removable space") + + for cache_key, entry in entries_to_evict: + del self._items[cache_key] + self._current_bytes -= entry.size_bytes + self._counters.evictions += len(entries_to_evict) + + def release(self, key: K) -> None: + """Release one acquired reference and apply retention policy. + + A RESERVED entry disappears when its final producer or follower + reference is released. A READY entry becomes evictable at zero + references and remains cached when `retain_after_release` is true; + otherwise it is removed immediately. + """ + with self._lock: + entry = self._items.get(key) + if entry is None or entry.reference_count == 0: + raise RuntimeError("cannot release an unreferenced cache entry") + + entry.reference_count -= 1 + if entry.reference_count: + return None + + if entry.state is _CacheEntryState.RESERVED: + self._reserved_bytes -= entry.size_bytes + del self._items[key] + return None + + if entry.state is not _CacheEntryState.READY: + raise RuntimeError(f"unexpected cache entry state: {entry.state}") + + self._in_use_bytes -= entry.size_bytes + if entry.retain_after_release: + return None + + self._current_bytes -= entry.size_bytes + del self._items[key] def pop(self, key: K) -> torch.Tensor | None: - """Remove one key and return its tensor, or `None` on miss.""" + """Remove one unreferenced key without applying release policy. + + Returns the READY tensor, or `None` for a removed reservation or miss. + """ with self._lock: - entry = self._items.pop(key, None) + entry = self._items.get(key) if entry is None: return None + if entry.reference_count: + raise RuntimeError("cannot remove a referenced cache entry") + + del self._items[key] + if entry.state is _CacheEntryState.RESERVED: + self._reserved_bytes -= entry.size_bytes + return None self._current_bytes -= entry.size_bytes self._prepare_for_current_stream(entry) + assert entry.value is not None return entry.value def clear(self) -> None: + """Remove all entries unless a managed reference or reservation is active.""" with self._lock: + if self._claimed_bytes: + raise RuntimeError("cannot clear cache with live references or reservations") self._items.clear() self._current_bytes = 0 @@ -209,6 +470,8 @@ def stats(self) -> TensorLRUCacheStats: return TensorLRUCacheStats( max_bytes=self._max_bytes, current_bytes=self._current_bytes, + reserved_bytes=self._reserved_bytes, + in_use_bytes=self._in_use_bytes, item_count=len(self._items), hits=self._counters.hits, misses=self._counters.misses, @@ -216,6 +479,8 @@ def stats(self) -> TensorLRUCacheStats: replacements=self._counters.replacements, evictions=self._counters.evictions, rejected_insertions=self._counters.rejected_insertions, + producer_misses=self._counters.producer_misses, + inflight_deduplications=self._counters.inflight_deduplications, hit_rate=self._counters.hit_rate, ) @@ -226,13 +491,30 @@ def log_stats(self, reason: str) -> None: f"bytes={stats.current_bytes}/{stats.max_bytes}, hits={stats.hits}, " f"misses={stats.misses}, hit_rate={stats.hit_rate:.3f}, " f"insertions={stats.insertions}, replacements={stats.replacements}, " - f"evictions={stats.evictions}, rejected_insertions={stats.rejected_insertions}" + f"evictions={stats.evictions}, rejected_insertions={stats.rejected_insertions}, " + f"producer_misses={stats.producer_misses}, " + f"inflight_deduplications={stats.inflight_deduplications}" ) + @property + def _claimed_bytes(self) -> int: + return self._reserved_bytes + self._in_use_bytes + @staticmethod def _tensor_size_bytes(tensor: torch.Tensor) -> int: return tensor.numel() * tensor.element_size() + def _clone_for_storage( + self, value: torch.Tensor + ) -> tuple[torch.Tensor, torch.cuda.Event | None]: + """Clone a value and record its producer event when stream-aware.""" + stored_value = value.detach().clone() + producer_event = None + if self._cuda_stream_aware and stored_value.is_cuda: + producer_event = torch.cuda.Event() + producer_event.record(torch.cuda.current_stream(stored_value.device)) + return stored_value, producer_event + def _prepare_for_current_stream(self, entry: _Entry) -> None: """Order and anchor a cached tensor for consumption on the current stream. @@ -243,6 +525,8 @@ def _prepare_for_current_stream(self, entry: _Entry) -> None: the storage while consumer-stream work is still pending, even if a later replacement or eviction drops the cache's own reference. """ + if entry.value is None: + raise RuntimeError("cannot access a cache entry before its value is stored") if not self._cuda_stream_aware or not entry.value.is_cuda: return @@ -251,13 +535,3 @@ def _prepare_for_current_stream(self, entry: _Entry) -> None: if entry.producer_event is not None: consumer_stream.wait_event(entry.producer_event) entry.value.record_stream(consumer_stream) - - def _evict_until_within_limit(self) -> tuple[int, int]: - evicted_count = 0 - evicted_bytes = 0 - while self._current_bytes > self._max_bytes: - _, entry = self._items.popitem(last=False) - self._current_bytes -= entry.size_bytes - evicted_count += 1 - evicted_bytes += entry.size_bytes - return evicted_count, evicted_bytes diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 243047f8113c..a68dd3f71ecc 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -631,11 +631,14 @@ class MultimodalConfig(StrictBaseModel): encoder_cache_max_bytes: NonNegativeInt = Field( default=134_217_728, # 128 MiB. description= - ("Maximum bytes for the opt-in multimodal encoder embedding cache; 0 " - "disables it. String values such as '512MB' and '1GiB' use binary " - "units. Inline encoding caches whole single-modality requests; item " - "scheduling caches individual items. Compatible with side-stream " - "prefetch; their memory limits are additive."), + ("Target capacity in bytes for reusable multimodal encoder embeddings; " + "0 disables cross-request reuse. With item scheduling, encoder outputs " + "still use the model-owned cache and its capacity is at least one legal " + "encoder iteration; this value can make that cache larger. String values " + "such as '512MB' and '1GiB' use binary units. Both inline encoding and " + "item scheduling cache individual items; inline caching supports only " + "single-modality parameters. " + "Compatible with side-stream prefetch."), status="prototype", ) diff --git a/tensorrt_llm/llmapi/rlhf_utils.py b/tensorrt_llm/llmapi/rlhf_utils.py index bf876058bf4b..247cb0cdd9d3 100644 --- a/tensorrt_llm/llmapi/rlhf_utils.py +++ b/tensorrt_llm/llmapi/rlhf_utils.py @@ -178,6 +178,13 @@ def update_weights(self, ipc_handles: Optional[dict] = None): weights[param_name] = tensor logger.info(f"weights key size: {len(weights.keys())}") + # Encoder outputs are weight-derived. The control action is + # draining by default, so any live reference or reservation + # here is a lifecycle error rather than state that may be + # force-dropped. Invalidate before every partial reload; a + # multi-call update session must not expose outputs from an + # earlier weight bucket through the persistent cache. + self.engine.invalidate_multimodal_encoder_cache() self.engine.model_engine.model_loader.reload( self.engine.model_engine.model, weights, allow_partial_loading=True ) diff --git a/tests/unittest/_torch/executor/engine/test_multimodal.py b/tests/unittest/_torch/executor/engine/test_multimodal.py index 4abd6fe31d5c..5fc2c3a5b560 100644 --- a/tests/unittest/_torch/executor/engine/test_multimodal.py +++ b/tests/unittest/_torch/executor/engine/test_multimodal.py @@ -1,9 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from collections.abc import Sequence from types import SimpleNamespace -from typing import Any import pytest import torch @@ -19,17 +17,16 @@ ) from tensorrt_llm._torch.pyexecutor.engine.multimodal import ( MultimodalItemScheduler, + resolve_bytes_per_mm_encoder_embedding, resolve_mm_encoder_output_budget, validate_mm_encoder_scheduling_compatibility, ) from tensorrt_llm._torch.pyexecutor.llm_request import ( - LlmRequest, MultimodalEncoderRequestError, initialize_multimodal_encoder_request, is_multimodal_encoder_ready, + make_mm_encoder_transient_cache_key, ) -from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache -from tensorrt_llm.bindings import SamplingConfig from tensorrt_llm.inputs.multimodal import MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalParams from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata from tensorrt_llm.llmapi.llm_args import MultimodalEncoderSchedulingPolicy @@ -38,71 +35,14 @@ pytestmark = pytest.mark.cpu_only -def _cache_request( - request_id: int, - *, - hashes: list[list[int]] | None, - embedding_lengths: Sequence[int], - kwargs_hash: str | None = "kw", -) -> LlmRequest: - """A cache-keyable item-scheduling request with raw image payload.""" - num_items = len(embedding_lengths) - multimodal_data = { - "image": { - "pixel_values": torch.arange(sum(embedding_lengths)).unsqueeze(1), - "image_grid_thw": torch.tensor([[1, 1, length] for length in embedding_lengths]), - }, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( - item_refs=[("image", item_idx) for item_idx in range(num_items)], - encoder_token_lengths=list(embedding_lengths), - output_embedding_lengths=list(embedding_lengths), - ), - "multimodal_embedding_lengths": list(embedding_lengths), - } - if kwargs_hash is not None: - multimodal_data["mm_processor_kwargs_hash"] = kwargs_hash - request = LlmRequest( - request_id=request_id, - max_new_tokens=1, - input_tokens=[1, 2, 3], - sampling_config=SamplingConfig(), - is_streaming=False, - py_multimodal_data=multimodal_data, - multimodal_hashes=hashes, - ) - initialize_multimodal_encoder_request(request, max_num_tokens=1 << 30) - return request - - -def _cache_mm_item_scheduler( - cache: TensorLRUCache[Any], - monkeypatch: pytest.MonkeyPatch, - *, - supports_encoder_cache: bool = True, -) -> MultimodalItemScheduler: - class _Model(MultimodalModelMixin): - supports_encoder_cache = False - - def __init__(self): - self.encoded_item_counts = [] - - def _get_multimodal_encoder_cache(self): - return cache - - def forward_multimodal_encoder_items(self, encoder_inputs): - # Items, not input tuples: adjacent same-request same-modality - # items are sliced into one tuple. - self.encoded_item_counts.append(sum(len(lengths) for _, lengths, _ in encoder_inputs)) - return [ - torch.full((embedding_length, 2), float(embedding_length)) - for _, embedding_lengths, _ in encoder_inputs - for embedding_length in embedding_lengths - ] - - monkeypatch.setattr(MultimodalParams, "to_device", lambda self, *args, **kwargs: self) - model = _Model() - model.supports_encoder_cache = supports_encoder_cache - return bare_mm_item_scheduler(model) +def _bind_items(mm_item_scheduler: MultimodalItemScheduler, request, *, row_bytes: int = 8) -> None: + state = request.py_mm_encoder_state + for item_idx, rows in enumerate(state.embedding_lengths): + cache_key = make_mm_encoder_transient_cache_key(request.request_id, item_idx) + assert mm_item_scheduler.encoder_cache.acquire( + cache_key, rows * row_bytes, retain_after_release=False + ) + state.set_item_cache_key(item_idx, cache_key, ready=False) def test_qwen3_output_budget_uses_post_merge_embedding_capacity() -> None: @@ -110,8 +50,6 @@ def test_qwen3_output_budget_uses_post_merge_embedding_capacity() -> None: processor = object.__new__(Qwen3VLInputProcessorBase) processor._config = SimpleNamespace(vision_config=SimpleNamespace(spatial_merge_size=2)) - # 16384 rows of fp16 -> 32768 bytes per embedding, via the mixin's explicit - # embedding_dim/embedding_dtype contract. model = SimpleNamespace(embedding_dim=16384, embedding_dtype=torch.float16) budget, bytes_per_embedding = resolve_mm_encoder_output_budget(processor, 65536, model) @@ -120,11 +58,18 @@ def test_qwen3_output_budget_uses_post_merge_embedding_capacity() -> None: assert budget == 512 * 1024**2 +def test_output_row_bytes_use_config_dtype_without_embedding_weight() -> None: + model = SimpleNamespace( + embedding_dim=16384, + model_config=SimpleNamespace(torch_dtype=torch.bfloat16), + ) + + assert resolve_bytes_per_mm_encoder_embedding(model) == 32768 + + def test_output_budget_requires_processor_embedding_capacity() -> None: processor = SimpleNamespace(get_max_mm_encoder_output_embeddings=lambda *_: None) - # The embedding capacity is validated before the model is consulted, so - # `model=None` never gets there. with pytest.raises(ValueError, match="get_max_mm_encoder_output_embeddings"): resolve_mm_encoder_output_budget(processor, 65536, None) @@ -135,6 +80,7 @@ def test_eager_compatibility_is_checked_only_for_item_scheduled_models() -> None encoder_scheduling_policy=MultimodalEncoderSchedulingPolicy.EAGER, encoder_side_stream_max_ahead=0, ), + pipeline_parallel_size=1, enable_attention_dp=True, cache_transceiver_config=SimpleNamespace(backend="NIXL"), ) @@ -155,6 +101,7 @@ def test_side_stream_compatibility_is_checked_only_for_item_scheduled_models() - encoder_scheduling_policy=MultimodalEncoderSchedulingPolicy.DEFAULT, encoder_side_stream_max_ahead=1, ), + pipeline_parallel_size=1, enable_attention_dp=False, cache_transceiver_config=None, ) @@ -165,11 +112,25 @@ def test_side_stream_compatibility_is_checked_only_for_item_scheduled_models() - validate_mm_encoder_scheduling_compatibility(args, item_scheduling_enabled=True) -def test_item_encoder_classifies_request_state_contract_errors() -> None: - class _Model(MultimodalModelMixin): - pass +def test_pipeline_parallel_compatibility_is_checked_only_for_item_scheduled_models() -> None: + args = SimpleNamespace( + multimodal_config=SimpleNamespace( + encoder_scheduling_policy=MultimodalEncoderSchedulingPolicy.DEFAULT, + encoder_side_stream_max_ahead=0, + ), + pipeline_parallel_size=2, + enable_attention_dp=False, + cache_transceiver_config=None, + ) - mm_item_scheduler = bare_mm_item_scheduler(_Model()) + validate_mm_encoder_scheduling_compatibility(args, item_scheduling_enabled=False) + + with pytest.raises(ValueError, match="pipeline parallelism"): + validate_mm_encoder_scheduling_compatibility(args, item_scheduling_enabled=True) + + +def test_item_encoder_classifies_request_state_contract_errors() -> None: + mm_item_scheduler = bare_mm_item_scheduler(MultimodalModelMixin()) request = make_llm_request(1) with pytest.raises(MultimodalEncoderRequestError, match="no longer active"): @@ -190,6 +151,7 @@ def forward_multimodal_encoder_items(self, _): mm_item_scheduler = bare_mm_item_scheduler(_Model()) request = make_mm_request(1, [4]) + _bind_items(mm_item_scheduler, request) with pytest.raises(MultimodalEncoderRequestError, match="one output per item"): mm_item_scheduler.forward_items([request], {request.request_id: [0]}) @@ -209,6 +171,7 @@ def forward_multimodal_encoder_items(self, _): mm_item_scheduler = bare_mm_item_scheduler(_Model()) request = make_mm_request(1, [4]) + _bind_items(mm_item_scheduler, request) expected = "bad request metadata" if failure_stage == "prepare" else "bad encoder output rows" with pytest.raises(MultimodalEncoderRequestError, match=expected): @@ -229,12 +192,13 @@ def forward_multimodal_encoder_items(self, _): mm_item_scheduler = bare_mm_item_scheduler(_Model()) request = make_mm_request(1, [4]) + _bind_items(mm_item_scheduler, request) with pytest.raises(torch.cuda.OutOfMemoryError, match="encoder OOM"): mm_item_scheduler.forward_items([request], {request.request_id: [0]}) -def test_item_outputs_accumulate_on_request_and_release_raw_data( +def test_item_outputs_commit_to_prompt_ordered_cache_keys( monkeypatch: pytest.MonkeyPatch, ) -> None: class _Model(MultimodalModelMixin): @@ -263,118 +227,26 @@ def forward_multimodal_encoder_items(self, encoder_inputs): }, ) initialize_multimodal_encoder_request(request, max_num_tokens=8) - assert request.py_mm_encoder_state.embedding_lengths == [2, 3] - - mm_item_scheduler.forward_items([request], {1: [0]}) - - # Items encoded across iterations land in their own rows of one buffer - # sized for the whole request, so the charge is already the full - # footprint. Raw inputs stay until every item is in. + _bind_items(mm_item_scheduler, request) state = request.py_mm_encoder_state - assert state.embeddings.shape == (2 + 3, 2) - assert state.recorded == [True, False] - assert state.resident_output_bytes(8) == (2 + 3) * 8 - assert "image" in request.py_multimodal_data - - mm_item_scheduler.forward_items([request], {1: [1]}) - - # Publishing is by reference: the buffer is already the contiguous form - # prefill consumes, so nothing is copied or concatenated here. - published = request.py_multimodal_data["multimodal_embedding"] - assert published is state.embeddings - assert published.tolist() == [ - [2.0, 2.0], - [2.0, 2.0], - [3.0, 3.0], - [3.0, 3.0], - [3.0, 3.0], - ] - assert "image" not in request.py_multimodal_data - assert state.resident_output_bytes(8) == (2 + 3) * 8 - assert is_multimodal_encoder_ready(request) - - -def test_duplicate_request_hits_cache_and_skips_encoding(monkeypatch: pytest.MonkeyPatch) -> None: - cache = TensorLRUCache(1 << 20, name="test") - mm_item_scheduler = _cache_mm_item_scheduler(cache, monkeypatch) - first = _cache_request(1, hashes=[[1, 2], [3, 4]], embedding_lengths=[2, 3]) - mm_item_scheduler.forward_items([first], {1: [0, 1]}) - assert mm_item_scheduler.model.encoded_item_counts == [2] - - second = _cache_request(2, hashes=[[1, 2], [3, 4]], embedding_lengths=[2, 3]) - mm_item_scheduler.forward_items([second], {2: [0, 1]}) - - # Read-through: every item hit, so the encoder never ran again, and each - # request owns an independent copy (no cross-request aliasing). - assert mm_item_scheduler.model.encoded_item_counts == [2] - assert is_multimodal_encoder_ready(second) - published = second.py_multimodal_data["multimodal_embedding"] - first_published = first.py_multimodal_data["multimodal_embedding"] - assert torch.equal(published, first_published) - assert published.untyped_storage().data_ptr() != first_published.untyped_storage().data_ptr() - - -def test_cache_hit_at_encode_skips_only_hit_items(monkeypatch: pytest.MonkeyPatch) -> None: - cache = TensorLRUCache(1 << 20, name="test") - mm_item_scheduler = _cache_mm_item_scheduler(cache, monkeypatch) - request = _cache_request(1, hashes=[[1, 2], [3, 4]], embedding_lengths=[2, 3]) - key0, _ = MultimodalModelMixin.build_encoder_cache_item_keys( - [[1, 2], [3, 4]], [("image", 0), ("image", 1)], [2, 3], "kw" - ) - cache.put(key0, torch.full((2, 2), 7.0)) # entry from an earlier request - mm_item_scheduler.forward_items([request], {1: [0, 1]}) - assert mm_item_scheduler.model.encoded_item_counts == [1] # only the miss encoded - published = request.py_multimodal_data["multimodal_embedding"] - assert torch.equal(published[:2], torch.full((2, 2), 7.0)) - assert is_multimodal_encoder_ready(request) - - -def test_cache_eviction_leaves_recorded_slots_intact(monkeypatch: pytest.MonkeyPatch) -> None: - cache = TensorLRUCache(2 * 2 * 4, name="test") # holds exactly one 2-row item - mm_item_scheduler = _cache_mm_item_scheduler(cache, monkeypatch) - request = _cache_request(1, hashes=[[1, 2]], embedding_lengths=[2]) - mm_item_scheduler.forward_items([request], {1: [0]}) - recorded = request.py_multimodal_data["multimodal_embedding"][:2] - assert len(cache) == 1 - - cache.clear() # simulate eviction of the entry the request came from + mm_item_scheduler.forward_items([request], {request.request_id: [0]}) - assert torch.equal(recorded, torch.full((2, 2), 2.0)) # owned clone, untouched - assert is_multimodal_encoder_ready(request) - - -def test_cache_off_encodes_every_item_without_touching_cache( - monkeypatch: pytest.MonkeyPatch, -) -> None: - # supports_encoder_cache=False -> mm_encoder_cache is None -> pure encode. - cache = TensorLRUCache(1 << 20, name="test") - mm_item_scheduler = _cache_mm_item_scheduler(cache, monkeypatch, supports_encoder_cache=False) - assert mm_item_scheduler.encoder_cache is None - request = _cache_request(1, hashes=[[1, 2], [3, 4]], embedding_lengths=[2, 3]) + assert state.item_ready == [True, False] + first = mm_item_scheduler.encoder_cache.get(state.item_cache_keys[0]) + torch.testing.assert_close(first, torch.full((2, 2), 2.0)) + assert "image" in request.py_multimodal_data - mm_item_scheduler.forward_items([request], {1: [0, 1]}) + mm_item_scheduler.forward_items([request], {request.request_id: [1]}) - assert mm_item_scheduler.model.encoded_item_counts == [2] + second = mm_item_scheduler.encoder_cache.get(state.item_cache_keys[1]) + torch.testing.assert_close(second, torch.full((3, 2), 3.0)) + assert "multimodal_embedding" not in request.py_multimodal_data + assert "image" not in request.py_multimodal_data assert is_multimodal_encoder_ready(request) - assert len(cache) == 0 # never populated - -@pytest.mark.parametrize("case", ["no_hashes", "no_kwargs_hash", "count_mismatch"]) -def test_key_guards_bypass_cache(case: str, monkeypatch: pytest.MonkeyPatch) -> None: - cache = TensorLRUCache(1 << 20, name="test") - mm_item_scheduler = _cache_mm_item_scheduler(cache, monkeypatch) - request = _cache_request( - 1, - hashes=None - if case == "no_hashes" - else ([[1, 2]] if case == "count_mismatch" else [[1, 2], [3, 4]]), - embedding_lengths=[2, 3], - kwargs_hash=None if case == "no_kwargs_hash" else "kw", + multimodal_data = mm_item_scheduler.build_multimodal_data_for_llm(request) + torch.testing.assert_close( + multimodal_data["multimodal_embedding"], + torch.cat([torch.full((2, 2), 2.0), torch.full((3, 2), 3.0)]), ) - - assert mm_item_scheduler.item_keys(request) is None - - mm_item_scheduler.forward_items([request], {1: [0, 1]}) - assert is_multimodal_encoder_ready(request) - assert len(cache) == 0 # unkeyable items never populate the cache diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 0d3a47329a7e..7547291ec378 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -24,6 +24,7 @@ from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MultimodalConfig, TorchLlmArgs from tensorrt_llm.mapping import Mapping @@ -66,6 +67,7 @@ class _EncoderCacheMultimodalModel(_MultimodalModel): class _ModelEngine: model: _TextModel | _MultimodalModel mm_encoder_output_budget_bytes: int | None = None + mm_encoder_cache: TensorLRUCache | None = None def _make_reserve_creator( @@ -83,9 +85,18 @@ def _make_reserve_creator( # Match ModelLoader: the model receives the effective, normalized config # from TorchLlmArgs rather than retaining the caller's input object. model.model_config.multimodal_config = llm_args.multimodal_config + cache_bytes = 0 + if isinstance(model, MultimodalModelMixin): + cache_bytes = mm_encoder_output_budget_bytes or 0 + if model.encoder_cache_active: + cache_bytes = max( + cache_bytes, + model.model_config.multimodal_config.encoder_cache_max_bytes, + ) model_engine = _ModelEngine( model=model, mm_encoder_output_budget_bytes=mm_encoder_output_budget_bytes, + mm_encoder_cache=TensorLRUCache(cache_bytes) if cache_bytes else None, ) return KvCacheCreator( model_engine=model_engine, @@ -387,6 +398,13 @@ def test_reserve_adds_only_unprofiled_output_capacity(): assert creator._get_multimodal_encoder_memory_reserve(profiled_output_bytes=400) == 112 +def test_downstream_pp_rank_without_encoder_store_reserves_no_memory() -> None: + creator = object.__new__(KvCacheCreator) + creator._model_engine = SimpleNamespace(mm_encoder_cache=None) + + assert creator._get_multimodal_encoder_memory_reserve() == 0 + + # --------------------------------------------------------------------------- # VSWA hybrid attention pool-group scaling (Gemma4 hybrid MMMU Pro hang fix) # --------------------------------------------------------------------------- @@ -947,6 +965,7 @@ def test_estimation_temporarily_uses_inferred_pool_sizing() -> None: # A bare Mock would auto-create the attribute; real engines set it to # None unless the model opted into MM item scheduling. model_engine.mm_encoder_output_budget_bytes = None + model_engine.mm_encoder_cache = None llm_args = Mock(cache_transceiver_config=None) with patch.object( diff --git a/tests/unittest/_torch/executor/multimodal_utils.py b/tests/unittest/_torch/executor/multimodal_utils.py index 6452721fe0da..fb4589cd9ee2 100644 --- a/tests/unittest/_torch/executor/multimodal_utils.py +++ b/tests/unittest/_torch/executor/multimodal_utils.py @@ -19,6 +19,7 @@ MultimodalEncoderRequestState, initialize_multimodal_encoder_request, ) +from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.bindings import SamplingConfig from tensorrt_llm.inputs.multimodal import MULTIMODAL_ENCODER_ITEM_METADATA_KEY from tensorrt_llm.inputs.registry import ( @@ -33,6 +34,8 @@ def bare_mm_item_scheduler( ) -> MultimodalItemScheduler: """A scheduler with no budget resolution -- the engine only builds one when item scheduling is engaged, so these tests skip `create()` and exercise the item path.""" + if model._multimodal_encoder_cache is None: + model._multimodal_encoder_cache = TensorLRUCache(1 << 20) return MultimodalItemScheduler(model=model, input_processor=input_processor) @@ -54,11 +57,9 @@ def record_output( hidden: int = 1, fill: float = 0.0, ) -> None: - """Write one item the way the encoder step does, sized from its declaration.""" - state.record( - item_idx, - torch.full((state.embedding_lengths[item_idx], hidden), fill), - ) + """Mark one item ready without materializing an output tensor.""" + del hidden, fill + state.item_ready[item_idx] = True def make_mm_request(request_id: int, costs: list[int], *, ready: Sequence[int] = ()) -> LlmRequest: diff --git a/tests/unittest/_torch/executor/test_multimodal_scheduler.py b/tests/unittest/_torch/executor/test_multimodal_scheduler.py index 3f98b02bf082..dee15b1e2458 100644 --- a/tests/unittest/_torch/executor/test_multimodal_scheduler.py +++ b/tests/unittest/_torch/executor/test_multimodal_scheduler.py @@ -32,6 +32,7 @@ ScheduledRequests, ) from tensorrt_llm._torch.pyexecutor.scheduler.waiting_queue import FCFSWaitingQueue +from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalParams, @@ -66,6 +67,30 @@ def can_schedule(self, requests): return bool(requests) +def _item_cache_keys(request): + state = request.py_mm_encoder_state + return [("test_mm", request.request_id, item_idx) for item_idx in range(state.num_items)] + + +def _scheduler( + *, + max_batch_size, + max_num_tokens, + cache_capacity=1 << 20, + base_scheduler=None, + scheduler_cls=MultimodalScheduler, +): + return scheduler_cls( + base_scheduler or _BaseScheduler(), + max_batch_size=max_batch_size, + max_num_tokens=max_num_tokens, + encoder_cache=TensorLRUCache(cache_capacity), + get_item_cache_keys=_item_cache_keys, + bytes_per_encoder_embedding=4, + retain_cache_entries=False, + ) + + def test_mm_encoder_token_lengths_distinguishes_missing_and_invalid_data(): request = make_llm_request(1) @@ -76,7 +101,7 @@ def test_mm_encoder_token_lengths_distinguishes_missing_and_invalid_data(): get_multimodal_encoder_token_lengths(request) -def test_mm_encoder_readiness_is_derived_from_request_local_outputs(): +def test_mm_encoder_readiness_is_derived_from_item_state(): request = make_mm_request(1, [4, 4]) assert request.py_mm_encoder_state.progress is MultimodalEncoderProgress.PENDING assert not is_multimodal_encoder_ready(request) @@ -106,7 +131,7 @@ def test_item_scheduling_rejects_raw_payload_without_item_metadata(): def test_multimodal_scheduler_keeps_items_atomic_and_backfills_requests(): - scheduler = MultimodalScheduler(_BaseScheduler(), max_batch_size=2, max_num_tokens=10) + scheduler = _scheduler(max_batch_size=2, max_num_tokens=10) first = make_mm_request(1, [7, 7]) second = make_mm_request(2, [3]) @@ -116,17 +141,33 @@ def test_multimodal_scheduler_keeps_items_atomic_and_backfills_requests(): assert output.context_requests == [second] -def test_scheduler_defers_items_beyond_output_byte_budget(): - # Budget hosts exactly one 1-row item (4 bytes): the second request's - # item must wait even though the token budget would admit it - # (allocate-before-compute). +def test_multimodal_scheduler_encodes_shared_cache_key_once(): + cache = TensorLRUCache(8) scheduler = MultimodalScheduler( _BaseScheduler(), - max_batch_size=8, - max_num_tokens=1 << 20, - output_budget_bytes=4, + max_batch_size=2, + max_num_tokens=8, + encoder_cache=cache, + get_item_cache_keys=lambda _request: [("stable", 0)], bytes_per_encoder_embedding=4, + retain_cache_entries=True, ) + first = make_mm_request(1, [4]) + second = make_mm_request(2, [4]) + + output = scheduler.schedule_request([first, second], set()) + + assert output.scheduled_mm_encoder_items == {first.request_id: [0]} + assert output.context_requests == [first, second] + assert first.py_mm_encoder_state.item_cache_keys == second.py_mm_encoder_state.item_cache_keys + assert cache.stats().inflight_deduplications == 1 + + +def test_scheduler_defers_items_beyond_output_byte_budget(): + # Budget hosts exactly one 1-row item (4 bytes): the second request's + # item must wait even though the token budget would admit it + # (allocate-before-compute). + scheduler = _scheduler(max_batch_size=8, max_num_tokens=1 << 20, cache_capacity=4) first = make_mm_request(1, [3]) second = make_mm_request(2, [3]) @@ -136,26 +177,25 @@ def test_scheduler_defers_items_beyond_output_byte_budget(): assert output.context_requests == [first] -def test_resident_outputs_of_active_requests_block_new_admissions(): - # A request that already holds recorded-but-unconsumed outputs (e.g. - # mid-chunked-prefill) occupies the budget purely through its live - # state — no counter, no release call — deferring new encoder work. - scheduler = MultimodalScheduler( - _BaseScheduler(), - max_batch_size=8, - max_num_tokens=1 << 20, - output_budget_bytes=4, - bytes_per_encoder_embedding=4, - ) - holder = make_mm_request(1, [3], ready=(0,)) # 1 row resident = full budget +def test_referenced_outputs_block_new_admissions_until_explicit_release(): + scheduler = _scheduler(max_batch_size=8, max_num_tokens=1 << 20, cache_capacity=4) + holder = make_mm_request(1, [3]) newcomer = make_mm_request(2, [3]) + first_output = scheduler.schedule_request([holder], set()) + assert first_output.scheduled_mm_encoder_items == {1: [0]} + holder_cache_key = holder.py_mm_encoder_state.item_cache_keys[0] + assert holder_cache_key is not None + scheduler.encoder_cache.commit(holder_cache_key, torch.ones(1, dtype=torch.float32)) + holder.py_mm_encoder_state.mark_cache_key_ready(holder_cache_key) + output = scheduler.schedule_request([holder, newcomer], set()) assert output.scheduled_mm_encoder_items is None - # Consumption (post-prefill strip clears the state) frees the budget on - # the next pass with no further bookkeeping — same for an aborted - # request, which simply leaves the active list. + drained = holder.py_mm_encoder_state.pop_all_cache_keys() + assert drained == [holder_cache_key] + scheduler.encoder_cache.release(holder_cache_key) + assert scheduler.encoder_cache.get(holder_cache_key, record_stats=False) is None holder.py_mm_encoder_state = None output = scheduler.schedule_request([holder, newcomer], set()) assert output.scheduled_mm_encoder_items == {2: [0]} @@ -166,13 +206,7 @@ def test_started_request_holds_its_whole_footprint_across_iterations(): # first item already allocates storage for all of them, so the bytes it # still needs are charged from the start. A request behind it cannot # squat that space and leave the head unable to finish. - scheduler = MultimodalScheduler( - _BaseScheduler(), - max_batch_size=8, - max_num_tokens=5, - output_budget_bytes=8, - bytes_per_encoder_embedding=4, - ) + scheduler = _scheduler(max_batch_size=8, max_num_tokens=5, cache_capacity=8) head = make_mm_request(1, [5, 5]) # second item exceeds this iteration's tokens follower = make_mm_request(2, [3]) @@ -213,34 +247,8 @@ def test_admission_rejects_requests_larger_than_output_budget(): assert "effective encoder_max_num_tokens is 1073741824" in str(exc_info.value) -def test_oversized_request_fails_fast_instead_of_starving(): - scheduler = MultimodalScheduler( - _BaseScheduler(), - max_batch_size=8, - max_num_tokens=1 << 20, - output_budget_bytes=4, - bytes_per_encoder_embedding=4, - ) - request = make_mm_request(1, [3, 3]) # 2 rows = 8 bytes > 4-byte budget - - with pytest.raises(RuntimeError, match="raise encoder_max_num_tokens") as exc_info: - scheduler.schedule_request([request], set()) - assert "Multimodal request 1" in str(exc_info.value) - assert "effective encoder_max_num_tokens is 1048576" in str(exc_info.value) - - -def test_scheduler_requires_bytes_per_embedding_alongside_budget(): - with pytest.raises(ValueError, match="bytes_per_encoder_embedding"): - MultimodalScheduler( - _BaseScheduler(), - max_batch_size=1, - max_num_tokens=1, - output_budget_bytes=4, - ) - - def test_multimodal_scheduler_selects_all_items_and_admits_request_when_batch_fits(): - scheduler = MultimodalScheduler(_BaseScheduler(), max_batch_size=2, max_num_tokens=10) + scheduler = _scheduler(max_batch_size=2, max_num_tokens=10) request = make_mm_request(1, [6, 4]) output = scheduler.schedule_request([request], set()) @@ -253,7 +261,7 @@ def test_multimodal_scheduler_selects_all_items_and_admits_request_when_batch_fi def test_multimodal_scheduler_respects_encoder_batch_size(): - scheduler = MultimodalScheduler(_BaseScheduler(), max_batch_size=2, max_num_tokens=4) + scheduler = _scheduler(max_batch_size=2, max_num_tokens=4) request = make_mm_request(1, [1, 1, 1, 1]) output = scheduler.schedule_request([request], set()) @@ -263,7 +271,7 @@ def test_multimodal_scheduler_respects_encoder_batch_size(): def test_multimodal_scheduler_withholds_request_on_budget_overflow(): - scheduler = MultimodalScheduler(_BaseScheduler(), max_batch_size=3, max_num_tokens=10) + scheduler = _scheduler(max_batch_size=3, max_num_tokens=10) request = make_mm_request(1, [6, 4, 1]) output = scheduler.schedule_request([request], set()) @@ -273,7 +281,7 @@ def test_multimodal_scheduler_withholds_request_on_budget_overflow(): def test_multimodal_scheduler_preserves_non_multimodal_requests(): - scheduler = MultimodalScheduler(_BaseScheduler(), max_batch_size=1, max_num_tokens=1) + scheduler = _scheduler(max_batch_size=1, max_num_tokens=1) request = make_llm_request(1) initialize_multimodal_encoder_request(request, max_num_tokens=1) @@ -294,7 +302,12 @@ def test_request_rejects_item_above_effective_startup_maximum(): def test_eager_scheduler_encodes_request_rejected_by_llm_capacity(): base_scheduler = _BaseScheduler() base_scheduler.capacity_scheduler = _RejectMultimodalCapacityScheduler() - scheduler = MultimodalEagerEncoderScheduler(base_scheduler, max_batch_size=1, max_num_tokens=8) + scheduler = _scheduler( + max_batch_size=1, + max_num_tokens=8, + base_scheduler=base_scheduler, + scheduler_cls=MultimodalEagerEncoderScheduler, + ) multimodal_request = make_mm_request(1, [8]) text_request = make_llm_request(2) initialize_multimodal_encoder_request(text_request, max_num_tokens=8) @@ -340,11 +353,12 @@ def test_forward_multimodal_encoder_step_contains_model_contract_error(): unrelated = make_llm_request(2) handled = [] - engine = SimpleNamespace( - forward_multimodal_encoder_items=bare_mm_item_scheduler( - MultimodalModelMixin() - ).forward_items - ) + def fail_encoder(*_): + raise MultimodalEncoderRequestError( + "multimodal_encoder_item_metadata must be a MultimodalEncoderItemMetadata" + ) + + engine = SimpleNamespace(forward_multimodal_encoder_items=fail_encoder) executor = object.__new__(PyExecutor) executor.active_requests = [failed, unrelated] @@ -602,20 +616,31 @@ def test_strip_mm_encoder_inputs_preserves_embedding_and_runtime_metadata(): assert "multimodal_embed_mask_cumsum" in mm_data -def test_terminate_request_releases_partial_multimodal_encoder_state(): +def test_terminate_request_releases_multimodal_cache_references_idempotently(): request = make_mm_request(1, [4, 4]) - record_output(request.py_mm_encoder_state, 0) + state = request.py_mm_encoder_state + cache = TensorLRUCache(16) + cache_key = ("mm_transient", request.request_id, 0) + cache.acquire(cache_key, 4, retain_after_release=False) + cache.commit(cache_key, torch.ones(1)) + state.set_item_cache_key(0, cache_key, ready=True) freed = [] executor = object.__new__(PyExecutor) + executor._mm_encoder_item_scheduling_enabled = True + executor.enable_attention_dp = False + executor.global_rank = 0 + executor.model_engine = SimpleNamespace(mm_encoder_cache=cache) executor.resource_manager = SimpleNamespace(free_resources=freed.append) executor._prefetched_request_ids = {request.py_request_id} executor._disagg_timed_out_ctx_cancelled_ids = {request.py_request_id} executor._disagg_timed_out_gen_cancelled_ids = {request.py_request_id} executor.gather_all_responses = False - executor.dist = SimpleNamespace(rank=1) + executor.dist = SimpleNamespace(rank=0) + executor.result_wait_queues = {} executor._do_terminate_request(request) + executor._release_multimodal_resources(request) assert freed == [request] assert request.py_mm_encoder_state is None @@ -625,9 +650,23 @@ def test_terminate_request_releases_partial_multimodal_encoder_state(): assert executor._disagg_timed_out_gen_cancelled_ids == set() -# --------------------------------------------------------------------------- -# Item-path read-through against the encoder cache (supports_encoder_cache) -# --------------------------------------------------------------------------- +def test_weight_invalidation_rejects_live_references(): + invalidations = [] + executor = object.__new__(PyExecutor) + executor.active_requests = [] + executor.model_engine = SimpleNamespace( + invalidate_multimodal_encoder_cache=lambda: invalidations.append(True) + ) + executor.invalidate_multimodal_encoder_cache() + + assert invalidations == [True] + + request = make_mm_request(1, [4]) + request.py_mm_encoder_state.set_item_cache_key(0, ("cache", 0), ready=False) + executor.active_requests = [request] + with pytest.raises(RuntimeError, match="live multimodal cache references"): + executor.invalidate_multimodal_encoder_cache() + assert invalidations == [True] # --------------------------------------------------------------------------- @@ -636,9 +675,9 @@ def test_terminate_request_releases_partial_multimodal_encoder_state(): def test_mm_encoder_state_enforces_lengths_slot_invariant(): - with pytest.raises(ValueError, match="one entry per item slot"): + with pytest.raises(ValueError, match="one cache key per item slot"): MultimodalEncoderRequestState( - embedding_lengths=[2], encoder_token_lengths=[4], recorded=[False, False] + embedding_lengths=[2], encoder_token_lengths=[4], item_ready=[False, False] ) @@ -651,111 +690,24 @@ def test_mm_encoder_state_copies_validated_scheduler_costs_at_admission(): metadata.encoder_token_lengths[0] = 100 assert request.py_mm_encoder_state.encoder_token_lengths == [4, 7] - scheduler = MultimodalScheduler(_BaseScheduler(), max_batch_size=2, max_num_tokens=11) + scheduler = _scheduler(max_batch_size=2, max_num_tokens=11) output = scheduler.schedule_request([request], set()) assert output.scheduled_mm_encoder_items == {1: [0, 1]} -def test_mm_encoder_state_progress_and_pending_transitions(): +def test_mm_encoder_state_tracks_prompt_ordered_cache_key_readiness(): state = MultimodalEncoderRequestState.from_embedding_lengths([2, 3]) + first_cache_key = ("cache", 0) + second_cache_key = ("cache", 1) - assert state.progress is MultimodalEncoderProgress.PENDING - assert state.pending_item_indices() == [0, 1] - - state.record(1, torch.ones(3, 2)) + state.set_item_cache_key(0, first_cache_key, ready=False) + state.set_item_cache_key(1, second_cache_key, ready=False) + state.mark_cache_key_ready(second_cache_key) assert state.progress is MultimodalEncoderProgress.PARTIAL assert state.pending_item_indices() == [0] - state.record(0, torch.zeros(2, 2)) + state.mark_cache_key_ready(first_cache_key) assert state.progress is MultimodalEncoderProgress.READY - # Items land in prompt order in their own row ranges of one buffer. - assert state.embeddings.tolist() == [ - [0, 0], - [0, 0], - [1, 1], - [1, 1], - [1, 1], - ] - - -def test_mm_encoder_state_record_rejects_mismatched_outputs(): - state = MultimodalEncoderRequestState.from_embedding_lengths([2, 3]) - - with pytest.raises(MultimodalEncoderRequestError, match="expected 2"): - state.record(0, torch.ones(5, 2)) - - state.record(0, torch.ones(2, 2)) - with pytest.raises(MultimodalEncoderRequestError, match="matching"): - state.record(1, torch.ones(3, 4)) # hidden dim mismatch vs items - with pytest.raises(MultimodalEncoderRequestError, match="already recorded"): - state.record(0, torch.ones(2, 2)) # items encode at most once - - -def test_mm_encoder_state_record_copies_into_owned_storage(): - state = MultimodalEncoderRequestState.from_embedding_lengths([2]) - batch = torch.arange(8, dtype=torch.float32).reshape(4, 2) - view = batch[1:3] # a view into a larger batched encoder output - - state.record(0, view) - - assert torch.equal(state.embeddings, view) - # The buffer neither aliases the batch storage (which a view would pin in - # full) nor can be invalidated by whoever owns the source tensor, and it - # is sized for this request alone. - assert state.embeddings.untyped_storage().data_ptr() != batch.untyped_storage().data_ptr() - assert state.embeddings.untyped_storage().nbytes() == 2 * 2 * 4 - - -def test_mm_encoder_state_charges_the_whole_request_from_its_first_item(): - state = MultimodalEncoderRequestState.from_embedding_lengths([2, 3]) - assert state.resident_output_bytes(4) == 0 - assert not state.has_storage - - # The first item allocates storage for every item, so the charge is the - # full footprint immediately -- which is what the request occupies. - state.record(1, torch.ones(3, 2)) - assert state.has_storage - assert state.resident_output_bytes(4) == (2 + 3) * 4 - - state.record(0, torch.ones(2, 2)) - assert state.resident_output_bytes(4) == (2 + 3) * 4 - - -def test_mm_encoder_state_finalize_is_a_conditional_no_op(): - state = MultimodalEncoderRequestState.from_embedding_lengths([2]) - multimodal_data = {"image": {"pixel_values": torch.empty(2, 1)}} - - assert state.finalize(multimodal_data) is False - assert "multimodal_embedding" not in multimodal_data - - state.record(0, torch.ones(2, 2)) - assert state.finalize(multimodal_data) is True - assert multimodal_data["multimodal_embedding"] is state.embeddings - assert "image" not in multimodal_data - - -def test_mm_encoder_state_publishes_the_buffer_without_copying(): - """Publishing must hand over the buffer itself, not a second materialization. - - The buffer is already the contiguous form the prefill path consumes, so a - copy here (or a per-item list that prefill has to concatenate) would put a - second full copy of the request's embeddings on the device that the byte - budget does not account for. - """ - state = MultimodalEncoderRequestState.from_embedding_lengths([2, 3]) - multimodal_data = {"image": {"pixel_values": torch.empty(2, 1)}} - state.record(0, torch.ones(2, 2)) - state.record(1, torch.ones(3, 2)) - buffer_ptr = state.embeddings.untyped_storage().data_ptr() - - assert state.finalize(multimodal_data) is True - - published = multimodal_data["multimodal_embedding"] - assert published is state.embeddings - assert published.untyped_storage().data_ptr() == buffer_ptr - assert published.shape == (2 + 3, 2) - # Readiness and byte accounting are unchanged by publishing: the rows stay - # resident until the request is stripped. - assert state.progress is MultimodalEncoderProgress.READY - assert state.pending_item_indices() == [] - assert state.resident_output_bytes(4) == (2 + 3) * 4 + assert state.pop_all_cache_keys() == [first_cache_key, second_cache_key] + assert state.item_cache_keys == [None, None] + assert state.progress is MultimodalEncoderProgress.PENDING diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index 9c4619413832..f1e8e587a5a8 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -147,6 +147,7 @@ def __init__(self, embedding_dim: int = 12) -> None: self.encoder_calls = 0 self.audio_tower = None self.embed_audio = object() + self._initialize_multimodal_encoder_cache(4096) @property def embedding_dim(self) -> int: diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py b/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py index 8ae672f430fe..75d2c530df06 100644 --- a/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py +++ b/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py @@ -74,6 +74,7 @@ def __init__(self, hidden_size: int, tokens_per_image: int): ) ) self.last_encoder_batch_size = 0 + self._initialize_multimodal_encoder_cache(4096) @property def embedding_dim(self) -> int: diff --git a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py index 1894a0c8551a..f019acec3edb 100644 --- a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py +++ b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py @@ -105,6 +105,7 @@ def __init__( multimodal_config=MultimodalConfig(encoder_cache_max_bytes=encoder_cache_max_bytes) ) self.encode_calls = 0 + self._initialize_multimodal_encoder_cache(encoder_cache_max_bytes) def encode_multimodal_inputs(self, multimodal_params, **encoder_kwargs) -> torch.Tensor: self.encode_calls += 1 @@ -294,15 +295,42 @@ def test_encoder_cache_requires_model_opt_in(): assert not model.encoder_cache_active -def test_encoder_cache_creation_logs_embedding_row_capacity(): +def test_explicit_cache_initialization_creates_cache_without_persistent_reuse(): + model = DummyMultimodalModel(make_embedding(hidden_size=4), torch.tensor([7])) + model.model_config = ModelConfig( + multimodal_config=MultimodalConfig(encoder_cache_max_bytes=4096) + ) + + cache = model._initialize_multimodal_encoder_cache(1024) + + assert cache is not None + assert cache.max_bytes == 1024 + assert model._multimodal_encoder_cache is cache + + +def test_explicit_cache_capacity_can_exceed_persistent_reuse_capacity(): model = CountingEncoderMultimodalModel( make_embedding(hidden_size=4), torch.tensor([7]), - encoder_cache_max_bytes=4096, + ) + model.model_config = ModelConfig( + multimodal_config=MultimodalConfig(encoder_cache_max_bytes=1024) ) + cache = model._initialize_multimodal_encoder_cache(2048) + + assert cache is not None + assert cache.max_bytes == 2048 + assert model._multimodal_encoder_cache is cache + + +def test_encoder_cache_creation_logs_embedding_row_capacity(): with patch("tensorrt_llm._torch.models.modeling_multimodal_mixin.logger.info") as info: - model._get_multimodal_encoder_cache() + CountingEncoderMultimodalModel( + make_embedding(hidden_size=4), + torch.tensor([7]), + encoder_cache_max_bytes=4096, + ) messages = [" ".join(map(str, call.args)) for call in info.call_args_list] assert any( @@ -319,7 +347,7 @@ def test_encoder_cache_creation_logs_byte_capacity_without_embedding_metadata(): ) with patch("tensorrt_llm._torch.models.modeling_multimodal_mixin.logger.info") as info: - model._get_multimodal_encoder_cache() + model._initialize_multimodal_encoder_cache(4096) messages = [" ".join(map(str, call.args)) for call in info.call_args_list] assert any( diff --git a/tests/unittest/_torch/test_tensor_lru_cache.py b/tests/unittest/_torch/test_tensor_lru_cache.py index cb65559394f3..b87c98d2f98c 100644 --- a/tests/unittest/_torch/test_tensor_lru_cache.py +++ b/tests/unittest/_torch/test_tensor_lru_cache.py @@ -19,7 +19,7 @@ import pytest import torch -from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache +from tensorrt_llm._torch.tensor_lru_cache import CacheAcquireResult, TensorLRUCache def test_rejects_non_positive_capacity() -> None: @@ -54,6 +54,22 @@ def test_put_get_pop_and_clear_update_byte_accounting() -> None: assert cache.get("a") is None +def test_clear_rejects_live_references_and_reservations() -> None: + cache = TensorLRUCache[str](max_bytes=32) + + assert cache.acquire("key", 8) is CacheAcquireResult.NEW_RESERVATION + with pytest.raises(RuntimeError, match="live references or reservations"): + cache.clear() + + cache.commit("key", torch.ones(2, dtype=torch.float32)) + with pytest.raises(RuntimeError, match="live references or reservations"): + cache.clear() + + assert cache.release("key") is None + cache.clear() + assert len(cache) == 0 + + def test_stats_track_hits_misses_insertions_and_replacements() -> None: cache = TensorLRUCache[str](max_bytes=32) tensor = torch.ones(2, dtype=torch.float32) @@ -182,6 +198,82 @@ def write_and_read(index: int) -> None: assert hit.numel() * hit.element_size() == 8 +def test_shared_reservation_stores_one_output_and_keeps_reusable_entry() -> None: + cache = TensorLRUCache[str](max_bytes=16) + value = torch.ones(2, dtype=torch.float32) + + assert cache.acquire("key", 8) is CacheAcquireResult.NEW_RESERVATION + assert cache.acquire("key", 8) is CacheAcquireResult.RESERVATION_HIT + assert cache.current_bytes == 0 + assert cache.stats().reserved_bytes == 8 + assert cache.stats().in_use_bytes == 0 + + cache.ensure_capacity(8) + cache.commit("key", value) + assert cache.current_bytes == 8 + assert cache.stats().reserved_bytes == 0 + assert cache.stats().in_use_bytes == 8 + cached = cache.get("key", record_stats=False) + assert cached is not None + torch.testing.assert_close(cached, value) + + assert cache.release("key") is None + assert cache.release("key") is None + assert cache.stats().in_use_bytes == 0 + assert cache.current_bytes == 8 + assert cache.acquire("key", 8) is CacheAcquireResult.READY_HIT + assert cache.release("key") is None + + stats = cache.stats() + assert stats.hits == 1 + assert stats.misses == 1 + assert stats.producer_misses == 1 + assert stats.inflight_deduplications == 1 + + +def test_non_retained_entries_are_removed_on_their_final_release() -> None: + cache = TensorLRUCache[str](max_bytes=16) + + assert ( + cache.acquire("reserved", 8, retain_after_release=False) + is CacheAcquireResult.NEW_RESERVATION + ) + assert cache.release("reserved") is None + assert len(cache) == 0 + assert cache.stats().reserved_bytes == 0 + + assert ( + cache.acquire("ready", 8, retain_after_release=False) is CacheAcquireResult.NEW_RESERVATION + ) + cache.commit("ready", torch.ones(2, dtype=torch.float32)) + assert cache.release("ready") is None + assert len(cache) == 0 + assert cache.current_bytes == 0 + assert cache.stats().in_use_bytes == 0 + + +def test_reservation_limit_and_output_space_are_checked_separately() -> None: + cache = TensorLRUCache[str](max_bytes=16) + assert cache.put("old-1", torch.ones(2, dtype=torch.float32)) + assert cache.put("old-2", torch.ones(2, dtype=torch.float32)) + + assert cache.acquire("new-1", 8) is CacheAcquireResult.NEW_RESERVATION + assert cache.acquire("new-2", 8) is CacheAcquireResult.NEW_RESERVATION + assert not cache.put("legacy", torch.ones(2, dtype=torch.float32)) + assert len(cache) == 4 + assert cache.current_bytes == 16 + assert cache.stats().reserved_bytes == 16 + assert cache.acquire("old-1", 8) is None + + cache.ensure_capacity(16) + assert cache.get("old-1", record_stats=False) is None + assert cache.get("old-2", record_stats=False) is None + assert cache.current_bytes == 0 + cache.commit("new-1", torch.ones(2, dtype=torch.float32)) + cache.commit("new-2", torch.ones(2, dtype=torch.float32)) + assert cache.current_bytes == 16 + + def test_stream_aware_mode_leaves_cpu_cache_behavior_unchanged() -> None: cache = TensorLRUCache[str](max_bytes=16, cuda_stream_aware=True) source = torch.arange(4, dtype=torch.float32) diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index 2a1731701e99..c1aa7bc3befa 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -12,6 +12,7 @@ MULTIMODAL_ENCODER_ITEM_METADATA_KEY, DisaggPrefillMultimodalInputs, MultimodalInput, + MultimodalParams, MultimodalRuntimeData, _find_mm_embedding_lengths_from_masks, find_mm_token_lengths, @@ -75,6 +76,13 @@ def test_mm_item_metadata_is_materialized_when_embedding_lengths_match(): assert "multimodal_item_refs" not in multimodal_data assert "multimodal_encoder_token_lengths" not in multimodal_data + params = MultimodalParams(multimodal_data=multimodal_data) + params.to_handle("multimodal_data") + assert isinstance( + params.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY], + MultimodalEncoderItemMetadata, + ) + def test_mm_item_metadata_rejects_mismatched_embedding_lengths(): input_processor = create_input_processor_with_hash(_ItemMetadataFakeProcessor([1]))