[pull] main from NVIDIA:main

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.

examples/models/core/deepseek_v3/README.md

@@ -863,10 +863,10 @@ echo "All processes completed!"
 The converted checkpoint could be used as `<YOUR_MODEL_DIR>` 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):
 

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

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

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(

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",

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 `<kAlignedBatchSize, split_kv, num_sms>`
+        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",