diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 46f90a9df31c..b03f9961f7cc 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -80,8 +80,8 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `Step3p7ForConditionalGeneration`| Yes | Yes | Yes | Untested | Untested | Yes | No | No | No | Yes | Untested | Untested | Yes | Untested | Untested | | `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes | Yes | Yes | Untested | Untested | No | No | No | No | Yes | Untested | No | N/A | Untested | Untested | -[^1]: Chunked Prefill for MLA can only be enabled on SM100/SM103. -[^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103 and in BF16/FP8 KV cache dtype. +[^1]: Chunked Prefill for MLA can only be enabled on SM90/SM100/SM103/SM120. +[^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103/SM120/SM121 and in BF16/FP8 KV cache dtype. [^3]: Qwen3-Next-80B-A3B exhibits relatively low accuracy on the SciCode-AA-v2 benchmark. [^5]: Supported via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend. See [AD Configs](../../../examples/auto_deploy/model_registry/configs). [^6]: Also supports text-only inference via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend. diff --git a/examples/models/core/deepseek_v3/README.md b/examples/models/core/deepseek_v3/README.md index 6c9774b83624..6f3af476866a 100644 --- a/examples/models/core/deepseek_v3/README.md +++ b/examples/models/core/deepseek_v3/README.md @@ -863,10 +863,10 @@ echo "All processes completed!" The converted checkpoint could be used as `` and consumed by other commands. ### KV Cache Reuse -KV cache reuse is supported for MLA on SM90, SM100 and SM120. It is enabled by default. Due to extra operations like memcpy and GEMMs, GPU memory consumption may be higher and the E2E performance may have regression in some cases. Users could pass `KvCacheConfig(enable_block_reuse=False)` to LLM API to disable it. +KV cache reuse is supported for MLA on SM90, SM100, SM103, SM120 and SM121. It is enabled by default. Due to extra operations like memcpy and GEMMs, GPU memory consumption may be higher and the E2E performance may have regression in some cases. Users could pass `KvCacheConfig(enable_block_reuse=False)` to LLM API to disable it. ### Chunked Prefill -Chunked Prefill is supported for MLA only on SM90 and SM100 currently. You should add `--enable_chunked_prefill` to enable it. The GPU memory consumption is highly correlated with `max_num_tokens` and `max_batch_size`. If encountering out-of-memory errors, you may make these values smaller. (`max_num_tokens` must be divisible by kv cache's `tokens_per_block`) +Chunked Prefill is supported for MLA on SM90, SM100, SM103 and SM120. You should add `--enable_chunked_prefill` to enable it. The GPU memory consumption is highly correlated with `max_num_tokens` and `max_batch_size`. If encountering out-of-memory errors, you may make these values smaller. (`max_num_tokens` must be divisible by kv cache's `tokens_per_block`) More specifically, we can imitate what we did in the [Quick Start](#quick-start): diff --git a/pyproject.toml b/pyproject.toml index 99b65e954d00..8110c926318e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ column_limit = 80 [tool.codespell] skip = ".git,3rdparty,triton_kernels,tests/integration/test_input_files**,**.jsonl,**.json" -ignore-words-list = "rouge,inout,atleast,strat,nd,subtile,thrid,improbe,NotIn,te,iteract,anythin,tru,Tracin,vEw,dOut,indext,asend,medias" +ignore-words-list = "rouge,inout,atleast,strat,nd,subtile,thrid,improbe,NotIn,te,iteract,anythin,tru,Tracin,vEw,dOut,indext,asend,medias,thw" [tool.autoflake] in-place = true diff --git a/requirements-dev.txt b/requirements-dev.txt index 8c6766a590eb..57d64fb0436c 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,4 +1,7 @@ -r requirements.txt +# VisualGen LPIPS goldens were recorded with this version. Install it before +# pytest collection so already-imported Diffusers modules match the package files. +diffusers==0.38.0 boto3 einops lpips diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 6d4c0a6249c1..2ae083d77355 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -1062,6 +1062,27 @@ def load_weights(self, weights: Dict[str, torch.Tensor]): def _parse_and_batch_multimodal_data( self, multimodal_params: List[MultimodalParams] ) -> Tuple[Dict[str, Any], Dict[str, List[Any]]]: + # Only the request itself knows how its image and video items + # interleave in the prompt, so a mixed request must carry a manifest. + for mp in multimodal_params: + data = mp.multimodal_data or {} + if ( + data.get("image") is not None + and data.get("video") is not None + and not mp.mm_item_order + ): + raise ValueError( + "Qwen3-VL mixed-modality requests must carry mm_item_order on MultimodalParams." + ) + + # Any batch that contains both modalities — whether within a single + # request or spread across heterogeneous single-modality requests — + # goes through the modality-blind ViT via one cat'd stream. + has_image = any(mp.multimodal_data.get("image") is not None for mp in multimodal_params) + has_video = any(mp.multimodal_data.get("video") is not None for mp in multimodal_params) + if has_image and has_video: + return self._interleave_multimodal_data(multimodal_params) + pixel_values_list = [] pixel_values_videos_list = [] image_grid_thw_list = [] @@ -1111,19 +1132,71 @@ def _parse_and_batch_multimodal_data( return mm_content_dict, mm_extra_data + def _interleave_multimodal_data( + self, multimodal_params: List[MultimodalParams] + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """Build one prompt-order pixel_values + grid_thw for mixed batches. + + The ViT is modality-blind (image = grid_thw row with t=1, video = t>1), + so we cat everything into a single stream in prompt order. Emits only + the "pixel_values" + "image_grid_thw" keys so forward() takes its + image branch uniformly over both modalities. + """ + pixels: List[torch.Tensor] = [] + grids: List[torch.Tensor] = [] + for mp in multimodal_params: + data = mp.multimodal_data + order = mp.mm_item_order or [] + img = data.get("image") or {} + vid = data.get("video") or {} + img_pv, img_thw = img.get("pixel_values"), img.get("image_grid_thw") + vid_pv = vid.get("pixel_values_videos") + vid_thw = vid.get("video_grid_thw") + + # Single-modality request in a mixed batch — no interleave + # needed, just append its rows. + if not order or img_pv is None or vid_pv is None: + if img_pv is not None: + pixels.append(img_pv) + grids.append(img_thw) + if vid_pv is not None: + pixels.append(vid_pv) + grids.append(vid_thw) + continue + + img_off = vid_off = 0 + for entry in order: + modality = entry["modality"] + idx = entry["index"] + if modality == "image": + thw = img_thw[idx] + n = int(thw.prod().item()) + pixels.append(img_pv[img_off : img_off + n]) + img_off += n + elif modality == "video": + thw = vid_thw[idx] + n = int(thw.prod().item()) + pixels.append(vid_pv[vid_off : vid_off + n]) + vid_off += n + else: + raise ValueError(f"Unknown modality in mm_item_order: {modality}") + grids.append(thw.unsqueeze(0)) + + return ( + {"pixel_values": torch.cat(pixels, dim=0)}, + {"image_grid_thw": torch.cat(grids, dim=0)}, + ) + @torch.inference_mode() def forward(self, multimodal_params: List[MultimodalParams]) -> List[torch.Tensor]: mm_content_data, mm_extra_data = self._parse_and_batch_multimodal_data(multimodal_params) pixel_values = mm_content_data.get("pixel_values", None) pixel_values_videos = mm_content_data.get("pixel_values_videos", None) - if pixel_values is not None and pixel_values_videos is not None: - raise ValueError("Currently only support single modality per request") - image_grid_thw = mm_extra_data.get("image_grid_thw", None) video_grid_thw = mm_extra_data.get("video_grid_thw", None) - embeds = [] + embeds: List[torch.Tensor] = [] if pixel_values is not None: pixel_values = pixel_values.to(self.model_dtype) image_embeds, deepstack_image_embeds = self.visual( diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 6d7522ec79a8..8ad581cf2070 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -677,6 +677,7 @@ def __init__( self.py_lora_path: str | None = kwargs.pop("py_lora_path", None) # Multimodal data self.py_multimodal_data = kwargs.pop("py_multimodal_data", None) + self.py_mm_item_order = kwargs.pop("py_mm_item_order", None) encoder_input_tokens = kwargs.get("encoder_input_tokens") encoder_output_len = kwargs.get("encoder_output_len") return_encoder_output = bool(kwargs.get("return_encoder_output", False)) @@ -1182,6 +1183,7 @@ def executor_request_to_llm_request( arrival_time=getattr(executor_request, "py_arrival_time", None), py_multimodal_data=getattr(executor_request, "py_multimodal_data", None), + py_mm_item_order=getattr(executor_request, "py_mm_item_order", None), kv_cache_retention_config=executor_request.kv_cache_retention_config, agent_hierarchy=agent_hierarchy, logprobs_mode=getattr(executor_request, "py_logprobs_mode", diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index d83f51a1bbe3..7aec823b110b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1159,6 +1159,17 @@ def warmup(self, resource_manager: ResourceManager) -> None: with self.cuda_graph_runner.allow_capture(): self._run_cuda_graph_warmup(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") + # Pre-compile DeepGEMM paged_mqa_logits_metadata for every 32-aligned + # batch bucket the runtime can produce (max_batch_size scaled by the + # MTP / DSL expansion factor when applicable). CUDA-graph warmup only + # exercises the batch sizes in cuda_graph_batch_sizes, which round + # up to a subset of buckets; any inference iter whose + # context_lens.size(0) lands on an uncovered bucket triggers an + # nvcc-driven JIT compile (~3s stall inside _prepare_inputs) on + # first touch. Pre-touching every bucket funnels that cost into + # warmup. No-op on non-DSA models. + self._warmup_dg_paged_mqa_logits_metadata() + log_mem_snapshot("warmup/after_dg_paged_mqa_logits_metadata") if can_run_general_warmup: # Pre-populate the memory pool with max-shape allocations to reduce # fragmentation at runtime. @@ -1167,6 +1178,102 @@ def warmup(self, resource_manager: ResourceManager) -> None: self._general_warmup(resource_manager, warmup_requests_configs) log_mem_snapshot("warmup/after_memory_pool_prepop") + def _warmup_dg_paged_mqa_logits_metadata(self) -> None: + """Pre-compile DeepGEMM's `get_paged_mqa_logits_metadata` helper for + every 32-aligned batch bucket the runtime can produce. + + DSA's `Indexer.prepare_scheduler_metadata` calls + `deep_gemm.get_paged_mqa_logits_metadata(context_lens, block_kv, + num_sms)` inside `_prepare_inputs` every iteration. The underlying + kernel is templated on `` + where `kAlignedBatchSize = align(context_lens.size(0), 32)` and + `split_kv` / `num_sms` are fixed for a given (block_kv, device). + deep_gemm's Python-side JIT compiles a fresh cubin (spawning + nvcc/cicc/ptxas, ~3s on GB300) the first time each `aligned_bs` + is requested. CUDA-graph warmup exercises only the batch sizes in + `cuda_graph_batch_sizes`, which round up to a subset of the 32- + aligned buckets; every uncovered bucket that the inference + workload later touches produces a 3s stall on that iteration. + Pre-touching every bucket here funnels those compiles into the + deterministic warmup phase. + + `context_lens.size(0)` is not always `num_generations`. For MTP + with `use_expanded_buffers_for_mtp=True` the expanded call passes + `num_generations * (1 + max_draft_tokens)`. For DSL expansion the + call passes `num_generations * dsl_expand_factor`, where + `dsl_expand_factor = next_n // eff` (`eff in kernel_atoms`, see + `_pick_dsl_expand` in `dsa.py`); its worst case is + `next_n = 1 + max_draft_tokens` when `eff == 1`. Reading the + current `dsl_expand_factor` off the metadata would under-estimate + the eventual max (it defaults to 1 before any prepare() has run, + and per-iter picks can differ across iters when CUDA graph is + off), so we use the static upper bound `1 + max_draft_tokens` + for both expansion paths. Bucket range is also scaled by + `max_beam_width` as a defense-in-depth ceiling for future beam + support (no-op today — DSA does not use beam). No-op on non-DSA + models. + + Best-effort: per-bucket JIT failures are logged and skipped so a + single broken bucket does not abort PyExecutor startup. + """ + attn_meta = getattr(self, "attn_metadata", None) + if attn_meta is None: + return + try: + from tensorrt_llm._torch.attention_backend.sparse.dsa import ( + _DG_SCHEDULE_BLOCK_KV, DSAtrtllmAttentionMetadata) + except ImportError: + return + if not isinstance(attn_meta, DSAtrtllmAttentionMetadata): + return + try: + from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata + except ImportError: + logger.info( + "[DG warmup] deep_gemm.get_paged_mqa_logits_metadata not " + "available; skipping paged_mqa_logits_metadata prewarm.") + return + + num_sms = attn_meta.num_sms + max_bs = max(1, int(self.batch_size)) + beam_width = max(1, int(getattr(self, "max_beam_width", 1) or 1)) + # Static upper bound on the row-count multiplier applied to + # `context_lens`. Both MTP-expanded and DSL-expanded call sites + # are bounded above by `(1 + max_draft_tokens)`; see the + # docstring for why we don't read the runtime `dsl_expand_factor` + # here. + max_draft_tokens = int(getattr(attn_meta, "max_draft_tokens", 0) or 0) + expands_batch = (getattr(attn_meta, "use_expanded_buffers_for_mtp", + False) + or getattr(attn_meta, "expand_for_dsl", False)) + expand_factor = 1 + max_draft_tokens if expands_batch else 1 + max_aligned = ((max_bs * beam_width * expand_factor + 31) // 32) * 32 + buckets = list(range(32, max_aligned + 32, 32)) + logger.info(f"[DG warmup] Pre-compiling paged_mqa_logits_metadata for " + f"{len(buckets)} aligned batch buckets up to {max_aligned} " + f"(block_kv={_DG_SCHEDULE_BLOCK_KV}, num_sms={num_sms}, " + f"max_bs={max_bs}, beam_width={beam_width}, " + f"expand_factor={expand_factor})") + for aligned_bs in buckets: + # Kernel scans `context_lens` and prefix-sums schedules; a + # zero-filled 2D tensor of shape (aligned_bs, 1) is enough to + # trigger dispatch and compile — the metadata output is + # discarded. + dummy = torch.zeros(aligned_bs, 1, dtype=torch.int32, device="cuda") + try: + _ = get_paged_mqa_logits_metadata(dummy, _DG_SCHEDULE_BLOCK_KV, + num_sms) + except RuntimeError as e: + # Narrow to RuntimeError so signature drifts in + # get_paged_mqa_logits_metadata (TypeError / ValueError) + # surface loudly instead of silently degrading perf. + logger.warning( + f"[DG warmup] paged_mqa_logits_metadata prewarm failed " + f"for aligned_bs={aligned_bs} " + f"(block_kv={_DG_SCHEDULE_BLOCK_KV}, num_sms={num_sms}); " + f"skipping bucket. {type(e).__name__}: {e}") + torch.cuda.synchronize() + def _general_warmup(self, resource_manager: ResourceManager, warmup_requests_configs: List[Tuple[int, int]]): """ @@ -3630,6 +3737,7 @@ def append_cross_attention_state(request: LlmRequest, request, self._mm_encoder_cache_enabled), multimodal_data=request.py_multimodal_data, multimodal_runtime=py_multimodal_runtime, + mm_item_order=getattr(request, "py_mm_item_order", None), input_ids_start_offset=context_start_idx) # Transfer any cross-iter MM encoder prefetch event stamped on the request onto the # freshly-built MultimodalParams. The downstream consume site reads it from the wrapper, @@ -4620,6 +4728,7 @@ def _prepare_tp_inputs_no_cache( multimodal_input=_build_request_multimodal_input( request, self._mm_encoder_cache_enabled), multimodal_data=request.py_multimodal_data, + mm_item_order=getattr(request, "py_mm_item_order", None), input_ids_start_offset=context_start_idx) multimodal_params.to_device("multimodal_data", "cuda", diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index a4c2e48ddc53..b984ec2a4024 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -48,6 +48,15 @@ from .model_loader import ModelLoader, _construct_checkpoint_loader from .py_executor import PyExecutor +_MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS = (90, 100, 103, 120, 121) +_MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS = (90, 100, 103, 120) +_MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS_STR = "/".join( + f"SM{sm_version}" + for sm_version in _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS) +_MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS_STR = "/".join( + f"SM{sm_version}" + for sm_version in _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS) + class _ExecutorMemoryMonitor: """Currently this focuses on tracking memory usage and related errors.""" @@ -704,12 +713,11 @@ def drafting_loop_wrapper(model): ) sm_version = get_sm_version() - if kv_cache_config.enable_block_reuse and sm_version not in [ - 90, 100, 103, 120 - ]: - logger.warning( - f"KV cache reuse for MLA can only be enabled on SM90/SM100/SM103/SM120, " - f"disable enable_block_reuse for SM{sm_version}") + if (kv_cache_config.enable_block_reuse and sm_version + not in _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS): + logger.warning("KV cache reuse for MLA can only be enabled on " + f"{_MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS_STR}, " + f"disable enable_block_reuse for SM{sm_version}") kv_cache_config.enable_block_reuse = False _set_model_engines_cache_reuse([model_engine, draft_model_engine], False) @@ -725,10 +733,11 @@ def drafting_loop_wrapper(model): kv_cache_config.enable_block_reuse = False _set_model_engines_cache_reuse([model_engine, draft_model_engine], False) - if enable_chunked_context and sm_version not in [90, 100, 103, 120]: - logger.warning( - "Chunked Prefill for MLA can only be enabled on SM90/SM100/SM103/SM120, " - f"disable enable_chunked_context for SM{sm_version}") + if (enable_chunked_context and sm_version + not in _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS): + logger.warning("Chunked Prefill for MLA can only be enabled on " + f"{_MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS_STR}, " + f"disable enable_chunked_context for SM{sm_version}") enable_chunked_context = False model_engine.attn_runtime_features.chunked_prefill = False if draft_model_engine is not None: diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 670d6f30508f..086ff957d54b 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -554,6 +554,8 @@ def _deduce_max_tokens(request: GenerationRequest, # E/P handoff embedding handles parked under "multimodal_embedding". request.multimodal_params.to_tensor("multimodal_data") executor_request.py_multimodal_data = request.multimodal_params.multimodal_data + if request.multimodal_params.mm_item_order: + executor_request.py_mm_item_order = request.multimodal_params.mm_item_order if self._is_pytorch_backend and request.sampling_params.logits_processor: # For PyTorch backend, we attach logits processors as a dynamic Python attribute diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 597011ed3fee..fb4b33779d32 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -3,7 +3,7 @@ """Multimodal utilities for handling images and other media types in TensorRT-LLM.""" from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Union import numpy as np import torch @@ -503,6 +503,12 @@ class MultimodalParams: multimodal_data: Optional[Dict[str, Any]] = field(default_factory=dict) multimodal_runtime: Optional[MultimodalRuntimeData] = None input_ids_start_offset: int = 0 + # Prompt-order manifest of data-backed items. Each entry is + # ``{"modality": m, "index": i, "placeholder": p}`` where ``i`` indexes + # ``multi_modal_data[m]``. Encoders that interleave items across modalities + # in prompt order read ``modality``/``index``; ``placeholder`` is carried + # for parity with the tracker manifest. + mm_item_order: Optional[List[Dict[str, Union[str, int]]]] = None # CUDA event recorded on a side stream by the MM encoder prefetch path. # When set, the consume site in `get_multimodal_embeddings` issues a # `wait_event` on the current stream before reading cached embeddings. @@ -814,11 +820,26 @@ def _update_hash(hasher, item: object) -> None: hasher.update(serialize_item(item)) -def apply_mm_hashes( - mm_data: Dict[str, Any], - mm_uuids: Optional[Dict[str, List[Optional[str]]]] = None, - hash_lib=default_hasher -) -> Tuple[Dict[str, List[str]], Optional[List[Optional[str]]]]: +class MMHashResult(NamedTuple): + """Parallel per-modality dicts of content hashes and their UUIDs. + + Both dicts are indexed first by modality name (`"image"`, `"video"`, + ...) then by within-modality item position, so `hashes[m][i]` and + `uuids[m][i]` describe the same item. + """ + + hashes: Dict[str, List[str]] + """Modality -> list of BLAKE3 hex digests (64 chars each).""" + + uuids: Optional[Dict[str, List[Optional[str]]]] + """Modality -> list of original UUID strings. `None` entries mark + items that fell back to content-only hashing. The whole dict is + `None` when the caller supplied no UUIDs at all.""" + + +def apply_mm_hashes(mm_data: Dict[str, Any], + mm_uuids: Optional[Dict[str, List[Optional[str]]]] = None, + hash_lib=default_hasher) -> "MMHashResult": """Apply hashing to multimodal data, one hash per multimodal item. When a UUID is provided for an item, the hash is computed from both the UUID @@ -834,9 +855,8 @@ def apply_mm_hashes( hash_lib: Hash function to use (default: blake3) Returns: - Tuple of: - - Dictionary of modality -> list of hash hex strings (64 chars each) - - Flattened list of original UUID strings (or None for content-hashed items) + `MMHashResult` with per-modality `hashes` and `uuids` dicts. See + the type's own docstring for the field-level contract. """ def _hash_item(item): @@ -868,9 +888,8 @@ def _hash_item_with_uuid(item, uuid: str): for modality, items in mm_data.items() } - # Collect UUIDs in the same order as items - all_uuids: List[Optional[str]] = [] mm_hashes: Dict[str, List[str]] = {} + mm_uuids_by_key: Dict[str, List[Optional[str]]] = {} for modality, items in mm_items.items(): modality_uuids = None @@ -884,22 +903,25 @@ def _hash_item_with_uuid(item, uuid: str): f"data items length ({len(items)}) for modality '{modality}'" ) - hashes = [] + hashes: List[str] = [] + uuids: List[Optional[str]] = [] for i, item in enumerate(items): uuid = modality_uuids[i] if modality_uuids else None if uuid is not None: # Hash UUID + content together for cache correctness hashes.append(_hash_item_with_uuid(item, uuid)) - all_uuids.append(uuid) # Store original UUID else: # Fall back to content-only hashing hashes.append(_hash_item(item)) - all_uuids.append(None) + uuids.append(uuid) # `None` when the caller didn't supply one mm_hashes[modality] = hashes + mm_uuids_by_key[modality] = uuids - # Return None for uuids if no UUIDs were provided at all - return mm_hashes, all_uuids if mm_uuids is not None else None + return MMHashResult( + hashes=mm_hashes, + uuids=mm_uuids_by_key if mm_uuids is not None else None, + ) def hexdigest_to_int32(hex_digest: str) -> List[int]: diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index e15e3e590314..6712c40196c4 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -1115,7 +1115,8 @@ def multimodal_hashing_process( # Extract optional UUIDs (can be None, or dict with same structure as mm_data) mm_uuids = inputs.get('multi_modal_uuids', None) - mm_hashes, mm_uuid_list = apply_mm_hashes(mm_data, mm_uuids, hash_lib) + mm_hashes, mm_uuids_by_key = apply_mm_hashes(mm_data, mm_uuids, + hash_lib) prompt_token_ids, extra_processed_inputs = input_processor( inputs, sampling_params) @@ -1147,7 +1148,17 @@ def multimodal_hashing_process( raise ValueError( "multimodal hashing could not determine multimodal token " "lengths for the provided input.") - num_mm_tokens = next(iter(num_mm_tokens_by_key.values())) + # `mm_hashes_flat`, `start_positions`, and `num_mm_tokens` must all + # index items at the same offsets — project into prompt order via the + # manifest. + mm_item_order = inputs.get("mm_item_order") + if mm_item_order: + num_mm_tokens = [ + num_mm_tokens_by_key[e["modality"]][e["index"]] + for e in mm_item_order + ] + else: + num_mm_tokens = next(iter(num_mm_tokens_by_key.values())) if len(num_mm_tokens) <= 0: raise ValueError("multimodal hashing produced an empty multimodal " "token-length list.") @@ -1193,8 +1204,30 @@ def multimodal_hashing_process( ) > 0 and mm_special_token_ids is not None: extra_processed_inputs["multimodal_data"][ "special_token_offsets"] = start_special_token_positions - # flatten the hashes from dict to a single list - mm_hashes_flat = [h for hashes in mm_hashes.values() for h in hashes] + # Same prompt-order projection as `num_mm_tokens` above — the cache + # key indexes each item's digest by its `start_positions` offset. + mm_item_order = inputs.get("mm_item_order") + if mm_item_order: + mm_hashes_flat = [ + mm_hashes[e["modality"]][e["index"]] for e in mm_item_order + ] + else: + mm_hashes_flat = [ + h for hashes in mm_hashes.values() for h in hashes + ] + # `MultimodalInput.multimodal_uuids` must index in lockstep with + # `multimodal_hashes`, so project through the same manifest. + if mm_uuids_by_key is None: + mm_uuid_list = None + elif mm_item_order: + mm_uuid_list = [ + mm_uuids_by_key[e["modality"]][e["index"]] + for e in mm_item_order + ] + else: + mm_uuid_list = [ + u for uuids in mm_uuids_by_key.values() for u in uuids + ] validate_mm_inputs(prompt_token_ids, mm_hashes_flat, start_positions, num_mm_tokens) mm_hashes_int32 = [hexdigest_to_int32(h) for h in mm_hashes_flat @@ -1211,19 +1244,22 @@ def process_prompt_maybe_hash( ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: try_multimodal_hashing = False # only used for first time use_multimodal_hashing = False # used for subsequent calls - modalities = list(set(inputs['multi_modal_data'].keys()) - ) if 'multi_modal_data' in inputs else [] - if len(modalities) > 0: - # TODO: support multimodal hashing for multiple modalities within the same request. - if len(modalities) == 1 and modalities[0] in [ - 'image', 'video', 'audio' - ]: - # only try multimodal hashing if the inputs only contain a single modality. - if input_processor.multimodal_hashing_supported is not None: - use_multimodal_hashing = input_processor.multimodal_hashing_supported - else: - # we need to try the multimodal hashing for the first time to determine if it is supported - try_multimodal_hashing = True + # Any subset of the supported modalities is eligible for hashing. + # ``None`` payloads are skipped so an empty bucket doesn't gate off + # hashing for the modalities that are actually present. + _SUPPORTED_HASHING_MODALITIES = ('image', 'video', 'audio') + modalities = [ + m for m, d in inputs.get('multi_modal_data', {}).items() + if d is not None + ] + if modalities and all(m in _SUPPORTED_HASHING_MODALITIES + for m in modalities): + if input_processor.multimodal_hashing_supported is not None: + use_multimodal_hashing = input_processor.multimodal_hashing_supported + else: + # First-time probe: attempt hashing and latch the result on + # the input processor. + try_multimodal_hashing = True if try_multimodal_hashing or use_multimodal_hashing: try: diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 7624256b223d..70ddbd8f055d 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -372,6 +372,14 @@ def __init__( self._embeddings = defaultdict[str, list](list) self._placeholder_counts = defaultdict[str, int](int) self._placeholder_to_modality: dict[str, str] = {} + # Prompt-order manifest of data-backed items. Each entry is + # `{"modality": , "index": , "placeholder": }`: + # `modality` is the modality name, `index` is the item's position in + # `multi_modal_data[modality]`, and `placeholder` is the exact string + # the input processor splices this item's embedding into. Populated by + # `add_data` (skipping `is_embedding=True` items — the interleave + # manifest addresses raw payload only). + self._item_order: list[dict[str, Union[str, int]]] = [] self._multimodal_server_config = multimodal_server_config if multimodal_server_config is not None else MultimodalServerConfig( ) # Per-request override merged with the server default at media-load @@ -440,6 +448,12 @@ def add_data(self, self._embeddings[media_type]) + 1 placeholder = retrieve_multimodal_placeholder(self._model_type, media_type, current_count) + if not is_embedding: + self._item_order.append({ + "modality": media_type, + "index": len(self._data[media_type]), + "placeholder": placeholder, + }) (self._embeddings if is_embedding else self._data)[media_type].append(data) if placeholder: @@ -455,20 +469,47 @@ def placeholder_modalities(self) -> Dict[str, str]: """Get the mapping from placeholder string to modality name.""" return dict(self._placeholder_to_modality) + def item_order(self) -> List[Dict[str, Union[str, int]]]: + """Prompt-order manifest of data-backed items. -def add_multimodal_placeholders(model_type: str, text_prompt: str, - mm_placeholder_counts: dict[str, int]) -> str: + Each entry is `{"modality": , "index": , "placeholder": }`. + `index` is the item's position in `multi_modal_data[modality]`; + `placeholder` is the exact string the input processor will look + for when splicing this item's encoder embedding. + """ + return list(self._item_order) + + +def add_multimodal_placeholders( + model_type: str, + text_prompt: str, + mm_placeholder_counts: dict[str, int], + item_order: Optional[List[Dict[str, Union[str, int]]]] = None, +) -> str: """Add multimodal placeholders to the text prompt. - Placeholders that already exist in the text are counted and subtracted - from the requested count to avoid double-insertion (e.g. when the - client already embeds ```` in the prompt text). + Placeholders already in the text are counted and subtracted to + avoid double-insertion (e.g. when the client already embeds + `` in the prompt text). When `item_order` is supplied, + placeholders are emitted in prompt-arrival order (needed for mixed + modality); otherwise they follow `mm_placeholder_counts` iteration + order. """ + if item_order: + wanted = [e["placeholder"] for e in item_order] + else: + wanted = [ + ph for ph, n in mm_placeholder_counts.items() for _ in range(n) + ] + + remaining = {ph: text_prompt.count(ph) for ph in set(wanted)} placeholders = [] - for placeholder, count in mm_placeholder_counts.items(): - existing = text_prompt.count(placeholder) - needed = max(0, count - existing) - placeholders.extend([placeholder] * needed) + for ph in wanted: + if remaining[ph] > 0: + remaining[ph] -= 1 + else: + placeholders.append(ph) + if not placeholders: return text_prompt parts = [] @@ -923,6 +964,9 @@ def convert_to_conversation_message( input[ "multi_modal_data"], _ = mm_data_tracker.retrieve_all_sync( ) + item_order = mm_data_tracker.item_order() + if item_order: + input["mm_item_order"] = item_order inputs.append(input) return inputs diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 5d19670b0991..3543414d1e01 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -890,12 +890,13 @@ def _preprocess( "prompt") # This is the text prompt, if present. if extra_processed_inputs is not None: query_token_ids = extra_processed_inputs.get('query_token_ids') - # Create unified MultimodalParams + # Create unified MultimodalParams. multimodal_params = MultimodalParams( multimodal_input=extra_processed_inputs.get( 'multimodal_input'), multimodal_data=extra_processed_inputs.get( - 'multimodal_data')) + 'multimodal_data'), + mm_item_order=inputs.get("mm_item_order")) # Only pass it if it has content if not multimodal_params.has_content(): multimodal_params = None diff --git a/tensorrt_llm/serve/chat_utils.py b/tensorrt_llm/serve/chat_utils.py index 023371adfe00..2ab0c8b362e7 100644 --- a/tensorrt_llm/serve/chat_utils.py +++ b/tensorrt_llm/serve/chat_utils.py @@ -482,6 +482,9 @@ def parse_chat_messages_coroutines( # Track placeholders added for this message only. msg_placeholder_counts = {} + # Snapshot the tracker's item_order length so we can slice off just + # the entries this message contributed. + item_order_start = len(mm_data_tracker.item_order()) if parsed_msg["media"]: for mdata in parsed_msg["media"]: placeholder = mm_data_tracker.add_data( @@ -506,13 +509,19 @@ def parse_chat_messages_coroutines( msg_placeholder_counts, mm_data_tracker.placeholder_modalities()) else: + msg_item_order = mm_data_tracker.item_order()[item_order_start:] parsed_msg["content"] = add_multimodal_placeholders( - type(model_config).model_type, parsed_msg["content"], - msg_placeholder_counts) + type(model_config).model_type, + parsed_msg["content"], + msg_placeholder_counts, + item_order=msg_item_order, + ) mm_placeholder_counts.append(msg_placeholder_counts) - return conversation, mm_data_tracker.retrieve_all_async( - ), mm_placeholder_counts + # ``item_order`` is populated synchronously by ``add_data``, so it can + # be returned directly (not through the coroutine). + return (conversation, mm_data_tracker.retrieve_all_async(), + mm_placeholder_counts, mm_data_tracker.item_order()) def make_tool_call_id(id_type: str = "random", func_name=None, idx=None): diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 53ab14af4234..3e1aa0d644db 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -1523,7 +1523,7 @@ async def chat_stream_generator( request.disaggregated_params) try: - conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines( + conversation, mm_coroutines, mm_placeholder_counts, mm_item_order = parse_chat_messages_coroutines( request.messages, self.model_config, self.multimodal_server_config, @@ -1532,7 +1532,7 @@ async def chat_stream_generator( # ValidatorIterator rejects extra fields; fall back to raw JSON. raw_body = await raw_request.json() raw_messages = raw_body.get("messages", []) - conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines( + conversation, mm_coroutines, mm_placeholder_counts, mm_item_order = parse_chat_messages_coroutines( raw_messages, self.model_config, self.multimodal_server_config, @@ -1574,6 +1574,8 @@ async def chat_stream_generator( raise ValueError( "Passing 'multi_modal_data' and 'multi_modal_embeddings' at the same time is not supported." ) + if mm_data and mm_item_order: + prompt["mm_item_order"] = mm_item_order if request.mm_processor_kwargs: prompt["mm_processor_kwargs"] = request.mm_processor_kwargs @@ -1709,7 +1711,7 @@ async def create_mm_embedding_response(promise: RequestOutput): ] try: - conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines( + conversation, mm_coroutines, mm_placeholder_counts, mm_item_order = parse_chat_messages_coroutines( request.messages, self.model_config, self.multimodal_server_config, @@ -1718,7 +1720,7 @@ async def create_mm_embedding_response(promise: RequestOutput): # ValidatorIterator rejects extra fields; fall back to raw JSON. raw_body = await raw_request.json() raw_messages = raw_body.get("messages", []) - conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines( + conversation, mm_coroutines, mm_placeholder_counts, mm_item_order = parse_chat_messages_coroutines( raw_messages, self.model_config, self.multimodal_server_config, @@ -1749,6 +1751,8 @@ async def create_mm_embedding_response(promise: RequestOutput): raise ValueError("Cannot use multimodal embeddings as input") if mm_data is not None: prompt["multi_modal_data"] = mm_data + if mm_item_order: + prompt["mm_item_order"] = mm_item_order promise = self.generator.generate_async(inputs=prompt, ) asyncio.create_task(self.await_disconnected(raw_request, promise)) diff --git a/tensorrt_llm/serve/resource_governor.py b/tensorrt_llm/serve/resource_governor.py index 20e5af208256..affd9b3a2ea0 100644 --- a/tensorrt_llm/serve/resource_governor.py +++ b/tensorrt_llm/serve/resource_governor.py @@ -120,7 +120,7 @@ async def _convert_messages( ) -> List[int]: """Convert chat messages to token IDs via chat template + tokenization.""" conversation: List[ConversationMessage] = [] - conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines( + conversation, mm_coroutines, mm_placeholder_counts, _ = parse_chat_messages_coroutines( messages, self.model_config, None ) token_task = async_apply_chat_template( diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index cfd363d71616..def6e4d907bb 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -820,7 +820,7 @@ async def _create_input_tokens( await conversation_store.store_messages(request.request_id, messages, request.previous_response_id) - conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines( + conversation, mm_coroutines, mm_placeholder_counts, _ = parse_chat_messages_coroutines( messages, model_config) tools_dict = [ tool.model_dump() diff --git a/tests/integration/defs/examples/visual_gen/conftest.py b/tests/integration/defs/examples/visual_gen/conftest.py new file mode 100644 index 000000000000..9d65d09ef545 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/conftest.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared fixtures for VisualGen example integration tests.""" + +import pytest +from defs.trt_test_alternative import check_call + + +@pytest.fixture(scope="session") +def _visual_gen_deps(llm_venv): + """Install optional media dependencies once per VisualGen test session.""" + llm_venv.run_cmd(["-m", "pip", "install", "av"]) + check_call(["apt-get", "update", "-y"], shell=False) + check_call(["apt-get", "install", "-y", "ffmpeg"], shell=False) diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip index b00403606a62..988c509f07b7 100644 --- a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:6d38286a6ee60db331a89bb3563fcc0896309b6e21228d4867ab33990c5b47ed -size 14386487 +oid sha256:78c7fde678df7b434b8f6da08d800d25bcfccb2509f73ab0c3d342eafe447100 +size 14648434 diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_t2v_fa4_fully_eager_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_t2v_fa4_fully_eager_lpips_golden_video.json new file mode 100644 index 000000000000..253c9f08e651 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_t2v_fa4_fully_eager_lpips_golden_video.json @@ -0,0 +1,25 @@ +{ + "video": "wan22_t2v_fa4_fully_eager_lpips_golden_video.mp4", + "model": "Wan2.2-T2V-A14B-Diffusers", + "source": "TensorRT-LLM VisualGen", + "prompt": "A cat sitting on a sunny windowsill watching birds outside.", + "negative_prompt": "", + "height": 480, + "width": 832, + "num_frames": 9, + "num_inference_steps": 4, + "guidance_scale": 4.0, + "seed": 42, + "frame_rate": 16.0, + "attention_backend": "FA4", + "torch_compile": false, + "torch_compile_stance": "force_eager", + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.25, + "diffusers_version": "0.38.0", + "torch_version": "2.12.0a0+0291f960b6.nv26.04.48445190", + "tensorrt_llm_version": "1.3.0rc21", + "tensorrt_llm_commit": "a0c406ff88c4a9736b5ce2f3c5eacbacdd0926d1", + "container_image": "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/release@sha256:0783d275ffe7efc3b6093dd6c7e823cd703eae973e462c8e1e4f7ca7dad60d64" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_t2v_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_t2v_lpips_golden_video.json index 35ecacdddc49..0621a4f06eec 100644 --- a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_t2v_lpips_golden_video.json +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/wan22_t2v_lpips_golden_video.json @@ -11,7 +11,9 @@ "guidance_scale": 4.0, "seed": 42, "frame_rate": 16.0, + "attention_backend": "VANILLA", "torch_compile": false, + "torch_compile_stance": "force_eager", "deterministic_algorithms": true, "lpips_net": "alex", "lpips_threshold": 0.05, diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen.py b/tests/integration/defs/examples/visual_gen/test_visual_gen.py index 83191a0d4837..ed8a5978941d 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen.py @@ -31,8 +31,6 @@ import pytest import torch import torch._inductor.config as inductor_config -from defs import conftest -from defs.common import venv_check_call from defs.trt_test_alternative import check_call from torch._inductor.async_compile import shutdown_compile_workers @@ -230,16 +228,6 @@ AESTHETIC_PREDICTOR_CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "emb_reader") -@pytest.fixture(scope="session") -def _visual_gen_deps(llm_venv): - """Install av + diffusers + ffmpeg once per session (shared by all video-gen fixtures).""" - llm_venv.run_cmd(["-m", "pip", "install", "av"]) - llm_venv.run_cmd(["-m", "pip", "install", "diffusers>=0.37.0"]) - # Install ffmpeg system package required by save_video() for MP4 encoding - check_call(["apt-get", "update", "-y"], shell=False) - check_call(["apt-get", "install", "-y", "ffmpeg"], shell=False) - - @pytest.fixture(scope="session") def vbench_repo_root(llm_venv): """Clone VBench repo into workspace and install; return repo root path.""" @@ -345,8 +333,30 @@ def _precache_aesthetic_predictor(): ) from exc +def _llm_models_root(): + # Imported lazily so that re-importing this module in a torch.multiprocessing.spawn + # child (a fresh interpreter) does not run a module-level `from defs import conftest`, + # which pulls in `tensorrt_llm.bindings` -- a compiled extension absent from the source + # tree the spawned child resolves, crashing the worker before the test runs. The parent + # process already imports conftest during collection, so this deferral is free. + from defs import conftest + + return conftest.llm_models_root() + + +def _venv_check_call(*args, **kwargs): + # Deferred like _llm_models_root above: defs.common does `from tensorrt_llm import + # LLM`, which pulls in tensorrt_llm.bindings. Importing it at module load would + # crash the torch.multiprocessing.spawn child processes used by the multi-GPU LPIPS + # tests, which re-import this module before the worker fixes sys.path. Only the + # single-GPU example tests call this, and only in the parent process. + from defs.common import venv_check_call + + return venv_check_call(*args, **kwargs) + + def _lpips_model_path(*parts): - return os.path.join(conftest.llm_models_root(), *parts) + return os.path.join(_llm_models_root(), *parts) def _skip_if_missing(path, label, is_dir=False): @@ -382,7 +392,7 @@ def _golden_media_path(tmp_path, media_name, label): def _ltx2_lpips_text_encoder_path(): - scratch_space = conftest.llm_models_root() + scratch_space = _llm_models_root() candidates = [ os.path.join(scratch_space, LTX2_TEXT_ENCODER_SUBPATH), os.path.join(scratch_space, "gemma", LTX2_TEXT_ENCODER_SUBPATH), @@ -412,14 +422,28 @@ def _cleanup_cuda(): @contextlib.contextmanager -def _lpips_deterministic_algorithms(): - previous = torch.are_deterministic_algorithms_enabled() - os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") - torch.use_deterministic_algorithms(True) +def _lpips_deterministic_algorithms(*, fully_eager=False): + previous_deterministic = torch.are_deterministic_algorithms_enabled() + previous_warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + previous_cublas_workspace_config = os.environ.get("CUBLAS_WORKSPACE_CONFIG") + try: - yield + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + torch.use_deterministic_algorithms(True) + compiler_context = ( + torch.compiler.set_stance("force_eager") if fully_eager else contextlib.nullcontext() + ) + with compiler_context: + yield finally: - torch.use_deterministic_algorithms(previous) + torch.use_deterministic_algorithms( + previous_deterministic, + warn_only=previous_warn_only, + ) + if previous_cublas_workspace_config is None: + os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None) + else: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = previous_cublas_workspace_config def _save_lpips_video_mp4(video, output_path, frame_rate): @@ -624,7 +648,7 @@ def _generate_ltx2_cuda_graph_trtllm_backend_video(output_path): TorchCompileConfig, ) - scratch_space = conftest.llm_models_root() + scratch_space = _llm_models_root() checkpoint_path = os.path.join(scratch_space, LTX2_MODEL_CHECKPOINT_PATH) text_encoder_path = _ltx2_lpips_text_encoder_path() spatial_upsampler_path = os.path.join(scratch_space, LTX2_UPSAMPLER_SUBPATH) @@ -699,6 +723,7 @@ def _run_wan_lpips_pipeline( seed, attention_backend="VANILLA", parallel=None, + fully_eager=False, ): from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig, VisualGenArgs @@ -712,7 +737,7 @@ def _run_wan_lpips_pipeline( ) if parallel is not None: args_kwargs["parallel_config"] = parallel - with _lpips_deterministic_algorithms(): + with _lpips_deterministic_algorithms(fully_eager=fully_eager): args = VisualGenArgs(**args_kwargs) pipeline = PipelineLoader(args).load(skip_warmup=True) try: @@ -747,7 +772,9 @@ def _generate_wan_lpips_video( guidance_scale, seed, frame_rate, + attention_backend="VANILLA", parallel=None, + fully_eager=False, ): generated_video = _run_wan_lpips_pipeline( model_path, @@ -759,7 +786,9 @@ def _generate_wan_lpips_video( num_inference_steps, guidance_scale, seed, + attention_backend=attention_backend, parallel=parallel, + fully_eager=fully_eager, ) assert generated_video is not None, "Single-GPU Wan LPIPS run produced no video" _save_lpips_video_mp4(generated_video, output_path, frame_rate=frame_rate) @@ -1151,7 +1180,7 @@ def _generate_wan_video(llm_venv, model_subpath, output_subdir): """ from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams - scratch_space = conftest.llm_models_root() + scratch_space = _llm_models_root() model_path = os.path.join(scratch_space, model_subpath) if not os.path.isdir(model_path): pytest.skip( @@ -1224,7 +1253,7 @@ def _generate_ltx2_two_stage_video(llm_venv, output_subdir, linear_type="default """ from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams - scratch_space = conftest.llm_models_root() + scratch_space = _llm_models_root() model_path = os.path.join(scratch_space, LTX2_MODEL_CHECKPOINT_PATH) text_encoder_path = os.path.join(scratch_space, LTX2_TEXT_ENCODER_SUBPATH) upsampler_path = os.path.join(scratch_space, LTX2_UPSAMPLER_SUBPATH) @@ -1384,7 +1413,7 @@ def _run_vbench_and_report( "custom_input", ] cmd.extend(["--dimension"] + VBENCH_DIMENSIONS) - venv_check_call(llm_venv, cmd) + _venv_check_call(llm_venv, cmd) pattern = os.path.join(output_path, "*_eval_results.json") result_files = glob.glob(pattern) @@ -1523,7 +1552,7 @@ def test_vbench_dimension_score_ltx2_two_stage_fp8( def test_visual_gen_quickstart(_visual_gen_deps, llm_root, llm_venv): """Run examples/visual_gen/quickstart_example.py end-to-end.""" - scratch_space = conftest.llm_models_root() + scratch_space = _llm_models_root() model_src = os.path.join(scratch_space, WAN_T2V_MODEL_SUBPATH) if not os.path.isdir(model_src): pytest.skip( @@ -1537,7 +1566,7 @@ def test_visual_gen_quickstart(_visual_gen_deps, llm_root, llm_venv): os.symlink(model_src, model_dst, target_is_directory=True) script_path = os.path.join(llm_root, "examples", "visual_gen", "quickstart_example.py") - venv_check_call(llm_venv, [script_path]) + _venv_check_call(llm_venv, [script_path]) output_path = os.path.join(llm_venv.get_working_directory(), "output.avi") assert os.path.isfile(output_path), f"Quickstart did not produce output.avi at {output_path}" @@ -1545,7 +1574,7 @@ def test_visual_gen_quickstart(_visual_gen_deps, llm_root, llm_venv): def test_visual_gen_api_walkthrough(_visual_gen_deps, llm_root, llm_venv): """Run examples/visual_gen/api_walkthrough.py end-to-end.""" - scratch_space = conftest.llm_models_root() + scratch_space = _llm_models_root() model_src = os.path.join(scratch_space, WAN_T2V_MODEL_SUBPATH) if not os.path.isdir(model_src): pytest.skip( @@ -1559,7 +1588,7 @@ def test_visual_gen_api_walkthrough(_visual_gen_deps, llm_root, llm_venv): os.symlink(model_src, model_dst, target_is_directory=True) script_path = os.path.join(llm_root, "examples", "visual_gen", "api_walkthrough.py") - venv_check_call(llm_venv, [script_path]) + _venv_check_call(llm_venv, [script_path]) output_path = os.path.join(llm_venv.get_working_directory(), "api_walkthrough_output.avi") assert os.path.isfile(output_path), f"API walkthrough did not produce {output_path}" @@ -1582,7 +1611,7 @@ def test_wan_t2v_example(_visual_gen_deps, llm_root, llm_venv): which runs the same script but with a no-quant YAML synthesized at runtime and additionally evaluates VBench scores. """ - scratch_space = conftest.llm_models_root() + scratch_space = _llm_models_root() model_path = os.path.join(scratch_space, WAN22_A14B_NVFP4_MODEL_SUBPATH) assert os.path.isdir(model_path), ( f"Model not found: {model_path} " @@ -1600,7 +1629,7 @@ def test_wan_t2v_example(_visual_gen_deps, llm_root, llm_venv): assert os.path.isfile(script_path), f"Example script not found: {script_path}" assert os.path.isfile(config_path), f"Config not found: {config_path}" - venv_check_call( + _venv_check_call( llm_venv, [ script_path, @@ -1636,7 +1665,7 @@ def test_flux1_example(_visual_gen_deps, llm_root, llm_venv): assert os.path.isfile(script_path), f"Example script not found: {script_path}" assert os.path.isfile(config_path), f"Config not found: {config_path}" - venv_check_call( + _venv_check_call( llm_venv, [ script_path, @@ -1672,7 +1701,7 @@ def test_flux2_example(_visual_gen_deps, llm_root, llm_venv): assert os.path.isfile(script_path), f"Example script not found: {script_path}" assert os.path.isfile(config_path), f"Config not found: {config_path}" - venv_check_call( + _venv_check_call( llm_venv, [ script_path, @@ -1711,7 +1740,7 @@ def test_ltx2_example(_visual_gen_deps, llm_root, llm_venv): assert os.path.isfile(script_path), f"Example script not found: {script_path}" assert os.path.isfile(config_path), f"Config not found: {config_path}" - venv_check_call( + _venv_check_call( llm_venv, [ script_path, @@ -1735,7 +1764,7 @@ def test_wan_i2v_example(_visual_gen_deps, llm_root, llm_venv): work together as documented. Uses the pre-quantized Wan 2.2 I2V A14B NVFP4 checkpoint and the default input image (cat_piano.png) bundled with the examples. """ - scratch_space = conftest.llm_models_root() + scratch_space = _llm_models_root() model_path = os.path.join(scratch_space, WAN22_I2V_A14B_NVFP4_MODEL_SUBPATH) if not os.path.isdir(model_path): pytest.skip( @@ -1754,7 +1783,7 @@ def test_wan_i2v_example(_visual_gen_deps, llm_root, llm_venv): assert os.path.isfile(script_path), f"Example script not found: {script_path}" assert os.path.isfile(config_path), f"Config not found: {config_path}" - venv_check_call( + _venv_check_call( llm_venv, [ script_path, @@ -1776,7 +1805,7 @@ def test_qwen_image_example(_visual_gen_deps, llm_root, llm_venv): ``configs/qwen-image-fp8-1gpu.yaml`` work together as documented. Uses the local Qwen-Image checkpoint and the shared FP8 blockwise dynamic-quant config. """ - scratch_space = conftest.llm_models_root() + scratch_space = _llm_models_root() model_path = os.path.join(scratch_space, QWEN_IMAGE_MODEL_SUBPATH) _skip_if_missing(model_path, "Qwen-Image checkpoint", is_dir=True) model_index_path = os.path.join(model_path, "model_index.json") @@ -1798,7 +1827,7 @@ def test_qwen_image_example(_visual_gen_deps, llm_root, llm_venv): assert os.path.isfile(script_path), f"Example script not found: {script_path}" assert os.path.isfile(config_path), f"Config not found: {config_path}" - venv_check_call( + _venv_check_call( llm_venv, [ script_path, @@ -1836,7 +1865,7 @@ def test_cosmos3_example(_visual_gen_deps, llm_root, llm_venv): assert os.path.isfile(script_path), f"Example script not found: {script_path}" assert os.path.isfile(config_path), f"Config not found: {config_path}" - venv_check_call( + _venv_check_call( llm_venv, [ script_path, diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index a60765fe600c..30c21d673f35 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -14,7 +14,9 @@ # limitations under the License. """Multi-GPU integration tests for VisualGen LPIPS quality checks.""" +import glob import os +import sys from typing import Callable import pytest @@ -39,16 +41,21 @@ _save_lpips_video_mp4, ) -try: - from tensorrt_llm._utils import get_free_port + +def _parallel_config(**kwargs): + # Imported lazily so that mp.spawn child processes resolve tensorrt_llm only after + # _distributed_worker has prepended the installed-wheel location to sys.path. A + # module-level tensorrt_llm import would run during the child's module re-import + # (before sys.path is fixed) and resolve to the bindings-less source tree. from tensorrt_llm.visual_gen.args import ParallelConfig - MODULES_AVAILABLE = True -except ImportError: - MODULES_AVAILABLE = False + return ParallelConfig(**kwargs) + # Keep it as 0.25 as the worst case scenario at NVL72 scale WAN_MULTI_GPU_LPIPS_THRESHOLD = 0.25 +WAN22_MULTI_GPU_LPIPS_ATTENTION_BACKEND = "FA4" +WAN22_MULTI_GPU_LPIPS_GOLDEN_VIDEO = "wan22_t2v_fa4_fully_eager_lpips_golden_video.mp4" WAN22_LPIPS_MULTI_GPU_VARIANTS = [ ("ulysses4", {"ulysses_size": 4}), ("cfg2_ulysses2", {"cfg_size": 2, "ulysses_size": 2}), @@ -90,7 +97,37 @@ def cleanup_distributed(): dist.destroy_process_group() -def _distributed_worker(rank, world_size, backend, test_fn, port, kwargs): +def _validated_tllm_site(site_dir): + """Return the realpath of ``site_dir`` after verifying it holds the installed wheel. + + The spawn workers rely on this directory to import tensorrt_llm with compiled + bindings; accepting an arbitrary path would let the import silently fall through + to the bindings-less source tree, so reject anything that does not contain the + package plus its compiled bindings extension. + """ + resolved = os.path.realpath(site_dir) if site_dir else "" + package_init = os.path.join(resolved, "tensorrt_llm", "__init__.py") + bindings = glob.glob(os.path.join(resolved, "tensorrt_llm", "bindings*.so")) + glob.glob( + os.path.join(resolved, "tensorrt_llm", "bindings", "*.so") + ) + if not (resolved and os.path.isfile(package_init) and bindings): + raise RuntimeError( + f"tllm_site={site_dir!r} does not contain an installed tensorrt_llm package " + "with compiled bindings; spawn workers would import the bindings-less " + "source tree instead of the wheel." + ) + return resolved + + +def _distributed_worker(rank, world_size, backend, test_fn, port, kwargs, tllm_site): + # mp.spawn starts a fresh interpreter whose sys.path (set up by the integration + # `defs` harness) puts the source checkout ahead of the installed wheel, so a bare + # `import tensorrt_llm` would resolve to the bindings-less source tree and crash the + # worker. Prepend the parent's installed-package location so the child imports + # tensorrt_llm (with compiled bindings) from the wheel before any such import. + tllm_site = _validated_tllm_site(tllm_site) + sys.path[:] = [path for path in sys.path if os.path.realpath(path) != tllm_site] + sys.path.insert(0, tllm_site) try: init_distributed_worker(rank, world_size, backend, port) test_fn(rank, world_size, **kwargs) @@ -102,22 +139,32 @@ def _distributed_worker(rank, world_size, backend, test_fn, port, kwargs): def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True, **kwargs): - if not MODULES_AVAILABLE: + try: + import tensorrt_llm.bindings as tllm_bindings + from tensorrt_llm._utils import get_free_port + except ImportError: pytest.skip("Required modules not available") if use_cuda and torch.cuda.device_count() < world_size: pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") backend = "nccl" if use_cuda else "gloo" port = get_free_port() + # Directory containing the installed tensorrt_llm package (i.e. site-packages), + # passed to spawn workers so they prepend it to sys.path and import the wheel with + # compiled bindings instead of the source-tree package. Validated here as well so a + # bad environment fails before any worker is spawned. + tllm_site = _validated_tllm_site( + os.path.dirname(os.path.dirname(os.path.abspath(tllm_bindings.__file__))) + ) mp.spawn( _distributed_worker, - args=(world_size, backend, test_fn, port, kwargs), + args=(world_size, backend, test_fn, port, kwargs, tllm_site), nprocs=world_size, join=True, ) def _skip_if_insufficient_gpus_for_parallel(parallel): - parallel_cfg = ParallelConfig(**parallel) + parallel_cfg = _parallel_config(**parallel) required = parallel_cfg.n_workers available = torch.cuda.device_count() if available < required: @@ -128,7 +175,7 @@ def _skip_if_insufficient_gpus_for_parallel(parallel): def _wan22_lpips_distributed_worker(rank: int, world_size: int, **kwargs) -> None: parallel = kwargs["parallel"] - ParallelConfig(**parallel).validate_world_size(world_size) + _parallel_config(**parallel).validate_world_size(world_size) generated_video = _run_wan_lpips_pipeline( kwargs["model_path"], @@ -140,8 +187,9 @@ def _wan22_lpips_distributed_worker(rank: int, world_size: int, **kwargs) -> Non kwargs["num_inference_steps"], kwargs["guidance_scale"], kwargs["seed"], - attention_backend="FA4", + attention_backend=WAN22_MULTI_GPU_LPIPS_ATTENTION_BACKEND, parallel=parallel, + fully_eager=True, ) if rank == 0: @@ -160,10 +208,12 @@ def _wan22_lpips_distributed_worker(rank: int, world_size: int, **kwargs) -> Non def _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel): _skip_if_insufficient_gpus_for_parallel(parallel) - parallel_cfg = ParallelConfig(**parallel) + parallel_cfg = _parallel_config(**parallel) generated_path = tmp_path / f"wan22_t2v_generated_{variant_name}.mp4" golden_path = _golden_media_path( - tmp_path, "wan22_t2v_lpips_golden_video.mp4", "Wan 2.2 LPIPS golden video" + tmp_path, + WAN22_MULTI_GPU_LPIPS_GOLDEN_VIDEO, + "Wan 2.2 FA4 fully-eager LPIPS golden video", ) run_test_in_distributed( @@ -200,7 +250,9 @@ def _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel): WAN22_LPIPS_MULTI_GPU_VARIANTS, ids=[name for name, _ in WAN22_LPIPS_MULTI_GPU_VARIANTS], ) -def test_wan22_t2v_lpips_against_golden_multi_gpu(tmp_path, variant_name, parallel): +def test_wan22_t2v_lpips_against_golden_multi_gpu( + _visual_gen_deps, tmp_path, variant_name, parallel +): _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel) @@ -209,5 +261,5 @@ def test_wan22_t2v_lpips_against_golden_multi_gpu(tmp_path, variant_name, parall WAN22_LPIPS_TP_VARIANTS, ids=[name for name, _ in WAN22_LPIPS_TP_VARIANTS], ) -def test_wan22_t2v_lpips_against_golden_tp(tmp_path, variant_name, parallel): +def test_wan22_t2v_lpips_against_golden_tp(_visual_gen_deps, tmp_path, variant_name, parallel): _run_wan22_t2v_lpips_case(tmp_path, variant_name, parallel) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6c107c0a219b..bd566aae915c 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -193,13 +193,6 @@ examples/test_whisper.py::test_llm_whisper_general[large-v3-disable_gemm_plugin- examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden SKIP (https://nvbugs/6418815) examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden SKIP (https://nvbugs/6437341) examples/visual_gen/test_visual_gen.py::test_ltx2_cuda_graph_trtllm_backend SKIP (https://nvbugs/6463822) -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[attn2d_2x2] SKIP (https://nvbugs/6272644) -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2] SKIP (https://nvbugs/6272644) -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[cfg2_ulysses2_attn2d_2x1] SKIP (https://nvbugs/6272644) -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_multi_gpu[ulysses4] SKIP (https://nvbugs/6272644) -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_tp[cfg2_tp2] SKIP (https://nvbugs/6329227) -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_tp[tp2] SKIP (https://nvbugs/6329227) -examples/visual_gen/test_visual_gen_multi_gpu.py::test_wan22_t2v_lpips_against_golden_tp[tp2_ulysses2] SKIP (https://nvbugs/6329227) full:A100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6422318) full:A100X/llmapi/test_llm_examples.py::test_llmapi_speculative_decoding_mtp SKIP (https://nvbugs/6287561) full:A100X/unittest/llmapi/test_llm_pytorch.py -m "part0" SKIP (https://nvbugs/6416249) diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index f55827f3e775..3b92fbfd2f25 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -15,7 +15,13 @@ from types import SimpleNamespace +import pytest + from tensorrt_llm._torch.pyexecutor import py_executor_creator +from tensorrt_llm._torch.pyexecutor.py_executor_creator import ( + _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS, + _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS, +) from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm.quantization import QuantAlgo @@ -184,8 +190,10 @@ def _make_llm_args(): ) -def _run_create_py_executor(monkeypatch, *, sm_version, kv_cache_quant_algo): - """Execute create_py_executor with mocked dependencies and return cache reuse flags. +def _run_create_py_executor( + monkeypatch, *, sm_version, kv_cache_quant_algo, enable_chunked_prefill=False +): + """Execute create_py_executor with mocked dependencies and return MLA runtime flags. Mocks all external dependencies (model engine, resource managers, etc.) to isolate executor creation logic and verify that KV cache reuse configuration is synchronized @@ -195,11 +203,14 @@ def _run_create_py_executor(monkeypatch, *, sm_version, kv_cache_quant_algo): monkeypatch: pytest fixture for mocking. sm_version: CUDA SM version to simulate (e.g., 89, 90). kv_cache_quant_algo: Quantization algorithm to use (e.g., NO_QUANT, INT8). + enable_chunked_prefill: Whether to request MLA chunked prefill support. Returns: - Tuple of (kv_cache_reuse_flag, runtime_cache_reuse_flag) from created executor. + Tuple of (kv_cache_reuse_flag, runtime_cache_reuse_flag, + runtime_chunked_prefill_flag) from created executor. """ llm_args = _make_llm_args() + llm_args.enable_chunked_prefill = enable_chunked_prefill fake_mapping = SimpleNamespace( rank=0, tp_size=1, @@ -275,6 +286,7 @@ def _create_py_executor_instance(**kwargs): return ( kv_cache_manager.enable_block_reuse, py_executor.model_engine.attn_runtime_features.cache_reuse, + py_executor.model_engine.attn_runtime_features.chunked_prefill, ) @@ -287,7 +299,7 @@ def test_mla_unsupported_sm_fallback_syncs_cache_reuse(monkeypatch): This test ensures invariant synchronization is maintained across the fallback. """ - kv_cache_reuse, runtime_cache_reuse = _run_create_py_executor( + kv_cache_reuse, runtime_cache_reuse, _ = _run_create_py_executor( monkeypatch, sm_version=89, kv_cache_quant_algo=QuantAlgo.NO_QUANT, @@ -307,7 +319,7 @@ def test_mla_unsupported_kv_quant_fallback_syncs_cache_reuse(monkeypatch): This test ensures invariant synchronization is maintained across the fallback. """ - kv_cache_reuse, runtime_cache_reuse = _run_create_py_executor( + kv_cache_reuse, runtime_cache_reuse, _ = _run_create_py_executor( monkeypatch, sm_version=90, kv_cache_quant_algo=QuantAlgo.INT8, @@ -317,21 +329,49 @@ def test_mla_unsupported_kv_quant_fallback_syncs_cache_reuse(monkeypatch): assert runtime_cache_reuse is False -def test_mla_supported_configuration_preserves_cache_reuse(monkeypatch): - """Verify MLA supported configuration preserves cache reuse in both config and runtime. +@pytest.mark.parametrize("sm_version", _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS) +def test_mla_supported_configuration_preserves_cache_reuse(monkeypatch, sm_version): + """Verify every supported MLA SM preserves cache reuse in both config and runtime. - When both SM version (90) and KV quantization (NO_QUANT) are supported for MLA, - no fallback occurs and: + When the SM version is in the MLA allowlist and KV quantization is + NO_QUANT, no unsupported-SM fallback should occur and: - kv_cache_config.enable_block_reuse remains True - model_engine.attn_runtime_features.cache_reuse remains True - - This positive test ensures the default path does not regress. """ - kv_cache_reuse, runtime_cache_reuse = _run_create_py_executor( + kv_cache_reuse, runtime_cache_reuse, _ = _run_create_py_executor( monkeypatch, - sm_version=90, + sm_version=sm_version, + kv_cache_quant_algo=QuantAlgo.NO_QUANT, + ) + + assert kv_cache_reuse is True + assert runtime_cache_reuse is True + + +@pytest.mark.parametrize("sm_version", _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS) +def test_mla_supported_configuration_preserves_chunked_prefill(monkeypatch, sm_version): + """Verify every supported MLA SM preserves chunked prefill when requested.""" + _, _, runtime_chunked_prefill = _run_create_py_executor( + monkeypatch, + sm_version=sm_version, + kv_cache_quant_algo=QuantAlgo.NO_QUANT, + enable_chunked_prefill=True, + ) + + assert runtime_chunked_prefill is True + + +def test_mla_sm121_fallback_preserves_cache_reuse_and_disables_chunked_prefill( + monkeypatch, +): + """Verify SM121 keeps cache reuse while disabling unsupported chunked prefill.""" + kv_cache_reuse, runtime_cache_reuse, runtime_chunked_prefill = _run_create_py_executor( + monkeypatch, + sm_version=121, kv_cache_quant_algo=QuantAlgo.NO_QUANT, + enable_chunked_prefill=True, ) assert kv_cache_reuse is True assert runtime_cache_reuse is True + assert runtime_chunked_prefill is False diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py index afedd960f9c3..426a6e659af8 100644 --- a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py +++ b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py @@ -1034,9 +1034,10 @@ def test_multi_request_batch_chat( inputs_with_embeddings): assert isinstance(input, dict) assert isinstance(input_with_embedding, dict) - assert list( - set(input.keys()) - ^ set(input_with_embedding.keys())) == ["multi_modal_data"] + assert (set(input.keys()) + ^ set(input_with_embedding.keys())) == { + "multi_modal_data", "mm_item_order" + } assert set(input_with_embedding.keys()) == set( ["prompt", "multi_modal_embeddings"]) assert input["prompt"] == input_with_embedding["prompt"] diff --git a/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py b/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py new file mode 100644 index 000000000000..1f0e01d900f6 --- /dev/null +++ b/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Encoder-level tests for Qwen3-VL mixed-modality interleaving. + +Exercises `Qwen3VisionModelBase._parse_and_batch_multimodal_data` and +`forward` directly by constructing an instance without invoking __init__ and +stubbing `self.visual` with a marker function whose output row values encode +input row identities. This lets us assert that mixed image+video requests +produce a single ViT call whose output rows arrive in prompt order. +""" + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VisionModelBase +from tensorrt_llm.inputs.multimodal import MultimodalParams + + +class _FakeVisual: + """Marker ViT that returns a single-tensor pair mimicking Qwen3VL's + (image_embeds, deepstack_list) contract. Row i of the output is a + length-`hidden` tensor filled with a marker computed from the input + row's marker (via pooling `patches_per_item` rows and preserving order). + """ + + def __init__(self, hidden=4, deepstack_layers=0): + self.hidden = hidden + self.deepstack_layers = deepstack_layers + + def __call__(self, pixel_values, grid_thw): + # `pixel_values` is [P_total, marker_dim]; each row's first column + # holds an integer identity marker. The "encoder" pools each item's + # patches into a single row per item (grid_thw specifies item extent). + out_rows = [] + offset = 0 + markers = pixel_values[:, 0].tolist() + for thw in grid_thw.tolist(): + n = int(thw[0] * thw[1] * thw[2]) + # Use the first patch's marker as the item's identity. + item_marker = markers[offset] + offset += n + # Emit ONE row per patch (real ViT output-length = P_total). + out_rows.extend([item_marker] * n) + base = torch.tensor(out_rows, dtype=torch.float32).unsqueeze(1).repeat(1, self.hidden) + deepstack = [torch.zeros_like(base) for _ in range(self.deepstack_layers)] + return base, deepstack + + +def _make_encoder(): + enc = object.__new__(Qwen3VisionModelBase) + enc.model_dtype = torch.float32 + enc.visual = _FakeVisual(hidden=4, deepstack_layers=0) + return enc + + +def _make_param(image=None, video=None, order=None): + # `image`, `video`, and `order` default to `None` so a test can omit + # whichever fields it doesn't need — every case populates at least one. + data = {} + if image is not None: + data["image"] = image + if video is not None: + data["video"] = video + return MultimodalParams(multimodal_data=data, mm_item_order=order) + + +def _make_item_patches(marker: int, patches: int, dim: int = 2): + # Row 0's first column encodes the item's marker. + t = torch.zeros((patches, dim), dtype=torch.float32) + t[:, 0] = float(marker) + return t + + +def test_image_only_single_modality_fallback_unchanged(): + enc = _make_encoder() + params = [ + _make_param( + image={ + "pixel_values": _make_item_patches(marker=10, patches=4), + "image_grid_thw": torch.tensor([[1, 2, 2]]), + }, + ) + ] + embeds = enc.forward(params) + assert len(embeds) == 1 + # 4 patches, all marker 10. + assert embeds[0].shape == (4, 4) + assert torch.equal(embeds[0][:, 0], torch.full((4,), 10.0)) + + +def test_mixed_image_video_image_prompt_order_interleave(): + enc = _make_encoder() + # Prompt order: image#0 (marker=10, 4 patches), video#0 (marker=20, 6 + # patches), image#1 (marker=30, 2 patches). + params = [ + _make_param( + image={ + "pixel_values": torch.cat( + [ + _make_item_patches(marker=10, patches=4), + _make_item_patches(marker=30, patches=2), + ], + dim=0, + ), + "image_grid_thw": torch.tensor([[1, 2, 2], [1, 1, 2]]), + }, + video={ + "pixel_values_videos": _make_item_patches(marker=20, patches=6), + "video_grid_thw": torch.tensor([[2, 1, 3]]), + }, + order=[ + {"modality": "image", "index": 0}, + {"modality": "video", "index": 0}, + {"modality": "image", "index": 1}, + ], + ) + ] + embeds = enc.forward(params) + assert len(embeds) == 1 + # 4 + 6 + 2 = 12 rows in prompt order. + expected = torch.tensor([10.0] * 4 + [20.0] * 6 + [30.0] * 2, dtype=torch.float32) + assert embeds[0].shape == (12, 4) + assert torch.equal(embeds[0][:, 0], expected) + + +def test_mixed_without_manifest_raises_tripwire(): + """Guard: if both modalities show up but mm_item_order is absent the + encoder must NOT silently produce misordered embeddings.""" + enc = _make_encoder() + params = [ + _make_param( + image={ + "pixel_values": _make_item_patches(marker=1, patches=1), + "image_grid_thw": torch.tensor([[1, 1, 1]]), + }, + video={ + "pixel_values_videos": _make_item_patches(marker=2, patches=1), + "video_grid_thw": torch.tensor([[1, 1, 1]]), + }, + # order deliberately omitted + ) + ] + with pytest.raises(ValueError, match="mm_item_order"): + enc.forward(params) + + +def test_batch_of_two_requests_single_modality_each(): + """Batching still works when each request in the batch is + single-modality (no manifest needed, no interleave path).""" + enc = _make_encoder() + params = [ + _make_param( + image={ + "pixel_values": _make_item_patches(marker=1, patches=2), + "image_grid_thw": torch.tensor([[1, 1, 2]]), + } + ), + _make_param( + image={ + "pixel_values": _make_item_patches(marker=2, patches=3), + "image_grid_thw": torch.tensor([[1, 1, 3]]), + } + ), + ] + embeds = enc.forward(params) + assert len(embeds) == 1 + expected = torch.tensor([1.0] * 2 + [2.0] * 3, dtype=torch.float32) + assert torch.equal(embeds[0][:, 0], expected) + + +def test_batch_unordered_mixed_request_raises_even_alongside_ordered_one(): + """A mixed-modality request without a manifest must still be rejected + per-request, even when a sibling request in the same batch carries a + valid manifest and would otherwise route the batch through the + interleave path.""" + enc = _make_encoder() + params = [ + _make_param( + image={ + "pixel_values": _make_item_patches(marker=10, patches=2), + "image_grid_thw": torch.tensor([[1, 1, 2]]), + }, + video={ + "pixel_values_videos": _make_item_patches(marker=20, patches=2), + "video_grid_thw": torch.tensor([[1, 1, 2]]), + }, + order=[ + {"modality": "image", "index": 0}, + {"modality": "video", "index": 0}, + ], + ), + _make_param( + image={ + "pixel_values": _make_item_patches(marker=30, patches=1), + "image_grid_thw": torch.tensor([[1, 1, 1]]), + }, + video={ + "pixel_values_videos": _make_item_patches(marker=40, patches=1), + "video_grid_thw": torch.tensor([[1, 1, 1]]), + }, + # No manifest on this request even though it has both modalities. + ), + ] + with pytest.raises(ValueError, match="mm_item_order"): + enc.forward(params) + + +def test_batch_of_image_only_and_video_only_requests_uses_interleave_path(): + """Two single-modality requests of different modalities in one batch + ride the modality-blind interleave path (one ViT call), even though + neither individual request is mixed and neither carries a manifest. + """ + enc = _make_encoder() + params = [ + _make_param( + image={ + "pixel_values": _make_item_patches(marker=1, patches=2), + "image_grid_thw": torch.tensor([[1, 1, 2]]), + } + ), + _make_param( + video={ + "pixel_values_videos": _make_item_patches(marker=2, patches=3), + "video_grid_thw": torch.tensor([[1, 1, 3]]), + } + ), + ] + embeds = enc.forward(params) + assert len(embeds) == 1 + expected = torch.tensor([1.0] * 2 + [2.0] * 3, dtype=torch.float32) + assert torch.equal(embeds[0][:, 0], expected) diff --git a/tests/unittest/inputs/test_chat_template_dispatch.py b/tests/unittest/inputs/test_chat_template_dispatch.py index 47627d27da4e..694b0c6f2163 100644 --- a/tests/unittest/inputs/test_chat_template_dispatch.py +++ b/tests/unittest/inputs/test_chat_template_dispatch.py @@ -381,7 +381,7 @@ async def fake_mm_coroutine(): monkeypatch.setattr( rg, "parse_chat_messages_coroutines", - lambda messages, model_config, _: ([], fake_mm_coroutine(), [{}]), + lambda messages, model_config, _: ([], fake_mm_coroutine(), [{}], None), ) # Must resolve the top-level model type, matching the serving call # sites (not the raw model_config.model_type). @@ -427,7 +427,7 @@ async def fake_mm_coroutine(): monkeypatch.setattr( ru, "parse_chat_messages_coroutines", - lambda messages, model_config: ([], fake_mm_coroutine(), [{}]), + lambda messages, model_config: ([], fake_mm_coroutine(), [{}], None), ) monkeypatch.setattr(ru, "resolve_top_level_model_type", lambda cfg: "resolved-model-type") monkeypatch.setattr(ru, "_get_chat_completion_function_tools", lambda tools: []) diff --git a/tests/unittest/llmapi/apps/test_chat_utils.py b/tests/unittest/llmapi/apps/test_chat_utils.py index 5c44c41770c1..cb401016d135 100644 --- a/tests/unittest/llmapi/apps/test_chat_utils.py +++ b/tests/unittest/llmapi/apps/test_chat_utils.py @@ -6,6 +6,7 @@ from tensorrt_llm.inputs.media_io import AudioMediaIO from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY +from tensorrt_llm.inputs.utils import retrieve_multimodal_placeholder from tensorrt_llm.serve import chat_utils as _chat_utils from tensorrt_llm.serve.chat_utils import ( _make_media_io, @@ -419,6 +420,20 @@ def test_jinjalike_literal(self): ) +def _expected_item(modality: str, index: int) -> dict: + """Build an expected ``item_order`` entry, mirroring ``add_data``. + + ``add_data`` records the placeholder produced for each item, formatting the + modality template with a 1-based running count that equals ``index + 1`` on + the data (non-embedding) path exercised by these tests. + """ + return { + "modality": modality, + "index": index, + "placeholder": retrieve_multimodal_placeholder(_MM_MODEL_TYPE, modality, index + 1), + } + + class TestMultimodalPlaceholderCounts: """Verify per-message multimodal placeholder counts. @@ -529,11 +544,78 @@ class _StubConfig: mock_config = _StubConfig() - _, _, mm_placeholder_counts = parse_chat_messages_coroutines(messages, mock_config, None) + _, _, mm_placeholder_counts, _ = parse_chat_messages_coroutines(messages, mock_config, None) assert mm_placeholder_counts == expected_mm_placeholder_counts +class TestMmItemOrderReturn: + """Tests for the ``mm_item_order`` return element. + + The 4th tuple element from ``parse_chat_messages_coroutines`` is the + ``MultimodalDataTracker.item_order()`` manifest, indexed within modality + in content-parts order. Each entry carries ``modality``, ``index`` and the + ``placeholder`` string recorded by ``add_data``. + """ + + @pytest.mark.parametrize( + "messages, expected", + [ + # Mixed image+video+image within one message: proves both + # content-parts order preservation (image, video, image) and + # per-modality index advance (image indices 0, 1; video 0). + ( + [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "a"}}, + {"type": "video_url", "video_url": {"url": "b"}}, + {"type": "image_url", "image_url": {"url": "c"}}, + ], + } + ], + [ + _expected_item("image", 0), + _expected_item("video", 0), + _expected_item("image", 1), + ], + ), + # Items spanning multiple messages: indices accumulate across + # messages, not per-message. + ( + [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "a"}}, + ], + }, + {"role": "assistant", "content": "ok"}, + { + "role": "user", + "content": [ + {"type": "video_url", "video_url": {"url": "b"}}, + {"type": "image_url", "image_url": {"url": "c"}}, + ], + }, + ], + [ + _expected_item("image", 0), + _expected_item("video", 0), + _expected_item("image", 1), + ], + ), + ], + ) + def test_item_order(self, messages, expected): + class _StubConfig: + model_type = _MM_MODEL_TYPE + + _, _, _, item_order = parse_chat_messages_coroutines(messages, _StubConfig(), None) + assert item_order == expected + + class TestParseChatMessageContentPart: """Unit tests for parse_chat_message_content_part.""" diff --git a/tests/unittest/llmapi/apps/test_chat_utils_validator_iterator.py b/tests/unittest/llmapi/apps/test_chat_utils_validator_iterator.py index 6b575ef849cc..1e248df5197f 100644 --- a/tests/unittest/llmapi/apps/test_chat_utils_validator_iterator.py +++ b/tests/unittest/llmapi/apps/test_chat_utils_validator_iterator.py @@ -151,7 +151,7 @@ def test_list_tool_calls(self): {"role": "assistant", "content": None, "tool_calls": [TOOL_CALL]}, {"role": "tool", "content": "72F", "tool_call_id": "call_1"}, ] - conv, _, _ = parse_chat_messages_coroutines(messages, self._mock_config(), None) + conv, _, _, _ = parse_chat_messages_coroutines(messages, self._mock_config(), None) assert len(conv) == 3 assert conv[1]["tool_calls"][0]["function"]["arguments"] == PARSED_ARGS @@ -160,7 +160,7 @@ def test_iterator_tool_calls(self): {"role": "user", "content": "hi"}, {"role": "assistant", "content": None, "tool_calls": SingleUseIterator([TOOL_CALL])}, ] - conv, _, _ = parse_chat_messages_coroutines(messages, self._mock_config(), None) + conv, _, _, _ = parse_chat_messages_coroutines(messages, self._mock_config(), None) assert conv[1]["tool_calls"][0]["function"]["arguments"] == PARSED_ARGS def test_extra_fields_raw_dict(self): @@ -169,7 +169,7 @@ def test_extra_fields_raw_dict(self): {"role": "user", "content": "hi"}, {"role": "assistant", "content": None, "tool_calls": [tc]}, ] - conv, _, _ = parse_chat_messages_coroutines(messages, self._mock_config(), None) + conv, _, _, _ = parse_chat_messages_coroutines(messages, self._mock_config(), None) assert conv[1]["tool_calls"][0]["name"] == "get_weather" diff --git a/tests/unittest/llmapi/test_llm_kv_cache_events.py b/tests/unittest/llmapi/test_llm_kv_cache_events.py index b170cd01219c..d2447c1aea2f 100644 --- a/tests/unittest/llmapi/test_llm_kv_cache_events.py +++ b/tests/unittest/llmapi/test_llm_kv_cache_events.py @@ -276,8 +276,8 @@ def test_apply_mm_hashes_with_uuids(): 0] # UUID changes hash # Second hash should be content-only (same as without UUID) assert hashes_partial["image"][1] == hashes_no_uuid["image"][1] - # UUIDs list should have the UUID and None - assert uuids_partial == ["sku-1234-a", None] + # UUIDs are returned per-modality, parallel to mm_hashes. + assert uuids_partial == {"image": ["sku-1234-a", None]} # Test with all UUIDs mm_uuids_all = {"image": ["sku-1234-a", "sku-1234-b"]} @@ -290,7 +290,7 @@ def test_apply_mm_hashes_with_uuids(): assert hashes_all["image"][1] != hashes_no_uuid["image"][1] # Different UUIDs with different content should produce different hashes assert hashes_all["image"][0] != hashes_all["image"][1] - assert uuids_all == ["sku-1234-a", "sku-1234-b"] + assert uuids_all == {"image": ["sku-1234-a", "sku-1234-b"]} def test_apply_mm_hashes_uuid_content_combined(): @@ -781,8 +781,9 @@ def test_apply_mm_hashes_multiple_modalities(): assert hashes["image"][0] != hashes_no_uuid["image"][0] assert hashes["video"][0] != hashes_no_uuid["video"][0] - # Check flattened UUID list (order may vary based on dict iteration) - assert set(uuids_list) == {"img-uuid-001", "vid-uuid-001"} + # UUIDs are returned per-modality; flatten across modalities to compare. + flat_uuids = [u for lst in uuids_list.values() for u in lst] + assert set(flat_uuids) == {"img-uuid-001", "vid-uuid-001"} def test_mm_keys_in_stored_events():