Mass-integration 0.20 to main

.devcontainer/docker-compose.yml

@@ -1,7 +1,7 @@
 version: "3.9"
 services:
   tensorrt_llm-dev:
-    image: urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-x86_64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505191345-4400
+    image: urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-x86_64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505211401-4539
     network_mode: host
     ipc: host
 

constraints.txt

@@ -1,2 +1,9 @@
+# These vulnerabilities were inherited from the base image (pytorch:25.05-py3) and should be removed when the base image
+# is updated.
+
 # WAR against https://github.com/advisories/GHSA-vqfr-h8mv-ghfj
 h11>=0.16.0
+# WAR against https://github.com/advisories/GHSA-7cx3-6m66-7c5m
+tornado>=6.5.0
+# WAR against https://github.com/advisories/GHSA-5rjg-fvgr-3xxf
+setuptools>=78.1.1

docker/Dockerfile.multi

@@ -72,9 +72,14 @@ RUN bash ./install_pytorch.sh $TORCH_INSTALL_TYPE && rm install_pytorch.sh
 RUN pip3 uninstall -y opencv && rm -rf /usr/local/lib/python3*/dist-packages/cv2/
 RUN pip3 install opencv-python-headless --force-reinstall --no-deps --no-cache-dir
 
-# WAR against https://github.com/advisories/GHSA-vqfr-h8mv-ghfj
-RUN pip3 install --upgrade h11>=0.16 --no-cache-dir
-
+# WARs against security issues inherited from pytorch:25.04
+# * https://github.com/advisories/GHSA-vqfr-h8mv-ghfj
+# * https://github.com/advisories/GHSA-7cx3-6m66-7c5m
+# * https://github.com/advisories/GHSA-5rjg-fvgr-3xxf
+RUN pip3 install --upgrade --no-cache-dir \
+    "h11>=0.16" \
+    "tornado>=6.5.0" \
+    "setuptools>=78.1.1,<80"
 
 FROM ${TRITON_IMAGE}:${TRITON_BASE_TAG} AS triton
 
@@ -173,5 +178,9 @@ RUN bash ./triton_backend/inflight_batcher_llm/scripts/build.sh
 FROM release AS tritonrelease
 
 WORKDIR /app/tensorrt_llm
-COPY ./triton_backend/ ./triton_backend/
+COPY ./triton_backend/all_models ./triton_backend/all_models
+COPY ./triton_backend/scripts ./triton_backend/scripts
+COPY ./triton_backend/tools ./triton_backend/tools
+COPY ./triton_backend/inflight_batcher_llm/scripts ./triton_backend/inflight_batcher_llm/scripts
+COPY ./triton_backend/inflight_batcher_llm/client ./triton_backend/inflight_batcher_llm/client
 COPY --from=tritonbuild /opt/tritonserver/backends/tensorrtllm /opt/tritonserver/backends/tensorrtllm

docs/source/advanced/kv-cache-management.md

@@ -0,0 +1,75 @@
+(kv-cache-management)=
+
+# KV Cache Management: Pools, Blocks, and Events
+
+This document provides an overview of the internal hierarchy and event system for paged KV cache management, as implemented in the TensorRT-LLM codebase.
+
+For more information on KV cache reuse see [KV cache reuse](kv-cache-reuse.md).
+
+---
+
+## Hierarchy: Pool, Block, and Page
+
+### **Block**
+- **Definition:** The smallest unit of KV cache allocation. A `KVCacheBlock` holds metadata (not the actual data) for a chunk of KV cache.
+- **Purpose:** Each block represents a fixed number of tokens' worth of KV data (can be specified by `tokens_per_block` parameter).
+- **Usage:** Blocks are allocated, reused, or evicted as sequences are processed.
+
+### **Page**
+- **Definition:** In this codebase, "page" is often used interchangeably with "block" (as in "paged KV cache"), but technically, a page could refer to a memory page (hardware-level), while a block is a logical unit for the cache.
+- **In Practice:** The code uses "block" as the main unit; "page" is not a distinct class or struct.
+
+### **Pool**
+- **Definition:** A pool is a contiguous memory buffer (or set of buffers) that holds the actual KV data for one or more layers.
+- **Types:** There are primary pools (fast GPU memory) and secondary pools (slower, e.g., CPU or offload memory).
+- **Organization:** Each pool can serve multiple layers that share the same KV head configuration. Pools are managed by `KVCacheBlockPool` and tracked in vectors in `WindowBlockManager`.
+- **Block ↔ Pool:** Each block is an index into a pool; the pool provides the actual storage, while the block is the metadata handle.
+
+### **WindowBlockManager/BlockManager**
+
+TRT-LLM supports 2 complex features related to KV cache management:
+1. **Variable Group-Query Attention (VGQA)** - i.e. a different `num_kv_heads` value for different layers.
+2. **Variable Sliding Window Attention (VSWA)** - i.e. a different `attention_window_size` value for different layers.
+
+In order to support both of these features, the pool management works as described below.
+
+But in the simple, *most common case*, for most models, where
+1. [MHA/MQA/Non-variable GQA](gpt-attention.md#multi-head-multi-query-and-group-query-attention), i.e., same `num_kv_heads` value for all layers,
+2. Global attention/[SWA](gpt-attention.md#sliding-window-attention-cyclic-rolling-buffer-kv-cache), i.e., same `attention_window_size` value for all layers,
+
+only a *single* pool will be created within the structure described below.
+
+#### KV Cache Pool Management
+
+- **WindowBlockManager:** Manages blocks and pools for a specific attention window size. Within a `WindowBlockManager`, there can be multiple pools - each corresponding a unique number of KV heads - i.e., to support VGQA.
+- **BlockManager:** Manages all `WindowBlockManager` instances, one per unique window size.
+
+**Hierarchy Summary:**
+- **Pool** (memory buffer for KV data)
+  - Contains many blocks.
+- **Blocks** (metadata for a chunk of the pool, each block = `tokens_per_block` tokens)
+    - (Optionally, blocks can be swapped between primary/secondary pools.)
+- **BlockManager/WindowBlockManager**: Manage pools and blocks, handle allocation, reuse, and eviction.
+
+---
+
+## Events in `KVCacheEventManager`
+
+The `KVCacheEventManager` is responsible for tracking and reporting significant changes in the state of the KV cache. Events are used for logging, debugging, or possibly for external monitoring.
+
+### **Types of Events**
+- **Created Event:** When pools or blocks are created/allocated.
+- **Updated Event:** When a block's state changes (e.g., moved between primary/secondary, priority updated).
+- **Removed Event:** When a block is removed from the cache (evicted or released).
+- **Stored Event:** When blocks are stored for potential reuse (e.g., after a sequence finishes and its blocks are reusable).
+
+### **What Triggers an Event?**
+- **Allocation/Deallocation:** Creating or freeing memory pools or blocks.
+- **Eviction/Reuse:** When a block is evicted, reused, or its priority changes.
+- **Block Movement:** When a block is moved between memory levels (primary ↔ secondary).
+- **Block Storage:** When blocks are stored for future reuse (e.g., after a sequence completes).
+
+**In summary:**
+An "event" is any significant change in the lifecycle or state of a KV cache block or pool, tracked for monitoring, debugging, or optimization purposes.
+
+---

docs/source/index.rst

@@ -104,6 +104,7 @@ Welcome to TensorRT-LLM's Documentation!
    advanced/inference-request.md
    advanced/lora.md
    advanced/expert-parallelism.md
+   advanced/kv-cache-management.md
    advanced/kv-cache-reuse.md
    advanced/speculative-decoding.md
    advanced/disaggregated-service.md

docs/source/torch/kv_cache_manager.md

@@ -4,6 +4,8 @@ In Transformer-based models, the KV (Key-Value) Cache is a mechanism used to opt
 Since KV Cache requires memory to store, it is also an important resource.
 In TensorRT-LLM, KV Cache is managed by the `KVCacheManager`.
 
+For details of the TensorRT-LLM `KVCacheManager` implementation see [KV Cache Management](../advanced/kv-cache-management.md).
+
 ## KV Cache Manager Introduction
 
 `KVCacheManager` is a type of resource manager, inheriting from `BaseResourceManager`.

examples/models/core/llama/README.md

@@ -128,7 +128,7 @@ trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp16_wq \
             --output_dir ./tmp/llama/7B/trt_engines/weight_only/1-gpu/ \
             --gemm_plugin auto
 
-# Build LLaMA 7B using 2-way auto parallelism.
+# Build LLaMA 7B using 2-way auto parallelism (deprecated).
 python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \
                             --output_dir ./tllm_checkpoint_1gpu_fp16 \
                             --dtype float16

examples/summarize.py

@@ -30,6 +30,9 @@
 import tensorrt_llm
 import tensorrt_llm.profiler as profiler
 from tensorrt_llm._utils import mpi_broadcast, str_dtype_to_torch
+from tensorrt_llm.builder import EngineConfig
+from tensorrt_llm.functional import RopeEmbeddingUtils, RotaryScalingType
+from tensorrt_llm.layers import MropeParams
 from tensorrt_llm.logger import logger
 from tensorrt_llm.models.qwen.utils import make_context
 from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelRunner
@@ -41,6 +44,42 @@
 from prompt_lookup.run_dtm_pld import run_dtm_pld
 
 
+def ensemble_mrope_params(batch_input_ids, max_position_embeddings,
+                          rotary_embedding_dim, theta):
+    mrope_params = MropeParams()
+    batch_size = len(batch_input_ids)
+
+    _, rotary_cos_sin = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin(
+        num_pos=max_position_embeddings,
+        dim=rotary_embedding_dim,
+        theta=1000000.0,
+        scale_type=RotaryScalingType.mrope,
+    )
+    rotary_cos_sin = torch.tensor(rotary_cos_sin).to(batch_input_ids[0].device)
+    rotary_cos_sin = rotary_cos_sin.reshape(max_position_embeddings,
+                                            int(rotary_embedding_dim / 2), 2)
+
+    cos_ori = rotary_cos_sin[:, :, 0]
+    sin_ori = rotary_cos_sin[:, :, 1]
+
+    mrope_position_ids_padding = torch.zeros(
+        (batch_size, max_position_embeddings), dtype=torch.int32)
+    for i in range(batch_size):
+        seq_len = batch_input_ids[i].shape[-1]
+        mrope_position_ids_padding[i, :seq_len] = torch.arange(
+            seq_len, device=batch_input_ids[i].device)
+
+    cos = cos_ori[mrope_position_ids_padding].unsqueeze(-1)
+    sin = sin_ori[mrope_position_ids_padding].unsqueeze(-1)
+
+    mrope_params.mrope_rotary_cos_sin = torch.concatenate(
+        (cos, sin), axis=-1).reshape(batch_size, -1)
+    mrope_params.mrope_position_deltas = torch.zeros(
+        [batch_size, 1], device=batch_input_ids[0].device)
+
+    return mrope_params
+
+
 def main(args):
     is_integration_test = os.getenv('INTEGRATION_TEST', '0') == '1'
     if is_integration_test:
@@ -262,7 +301,19 @@ def eval_trt_llm(datapoint,
                                           eval_task=eval_task,
                                           add_special_tokens=add_special_tokens,
                                           min_input_length=min_input_length)
-        batch_size = len(batch_input_ids)
+        # Generate mrope params for qwen model
+        engine_config = EngineConfig.from_json_file(
+            f"{args.engine_dir}/config.json")
+        pretrain_config = engine_config.pretrained_config
+        mrope_params = None
+        if 'qwen' in model_name.lower():
+            mrope_params = ensemble_mrope_params(
+                batch_input_ids,
+                max_position_embeddings=pretrain_config.max_position_embeddings,
+                rotary_embedding_dim=pretrain_config.rotary_embedding_dim,
+                theta=pretrain_config.rotary_base,
+            )
+
         if batch_size == 0:
             return [], [], [], {}
         input_lengths = [x.size(0) for x in batch_input_ids]
@@ -309,7 +360,8 @@ def eval_trt_llm(datapoint,
                     return_dict=True,
                     random_seed=random_seed,
                     medusa_choices=args.medusa_choices,
-                    eagle_choices=args.eagle_choices)
+                    eagle_choices=args.eagle_choices,
+                    mrope_params=mrope_params)
                 torch.cuda.synchronize()
 
         # Extract a list of tensors of shape beam_width x output_ids.

jenkins/L0_MergeRequest.groovy

@@ -28,10 +28,10 @@ UPLOAD_PATH = env.uploadPath ? env.uploadPath : "sw-tensorrt-generic/llm-artifac
 // Container configuration
 // available tags can be found in: https://urm.nvidia.com/artifactory/sw-tensorrt-docker/tensorrt-llm/
 // [base_image_name]-[arch]-[os](-[python_version])-[trt_version]-[torch_install_type]-[stage]-[date]-[mr_id]
-LLM_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-x86_64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505191345-4400"
-LLM_SBSA_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-aarch64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505191345-4400"
-LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py310-trt10.10.0.31-skip-tritondevel-202505191345-4400"
-LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py312-trt10.10.0.31-skip-tritondevel-202505191345-4400"
+LLM_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-x86_64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505211401-4539"
+LLM_SBSA_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-aarch64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505211401-4539"
+LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py310-trt10.10.0.31-skip-tritondevel-202505211401-4539"
+LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py312-trt10.10.0.31-skip-tritondevel-202505211401-4539"
 
 // TODO: Move common variables to an unified location
 BUILD_CORES_REQUEST = "8"

jenkins/controlCCache.groovy

@@ -1,7 +1,7 @@
 
 import java.lang.InterruptedException
 
-DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-x86_64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505191345-4400"
+DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.04-py3-x86_64-ubuntu24.04-trt10.10.0.31-skip-tritondevel-202505211401-4539"
 
 def createKubernetesPodConfig(image)
 {

tensorrt_llm/_torch/distributed/ops.py

@@ -1,6 +1,7 @@
 import math
 import os
 import threading
+from itertools import accumulate
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -116,6 +117,24 @@ def get_output_info(input: torch.Tensor, dim: int) -> List[int]:
     return {'output_shape': output_shape, 'numel_base': numel_base}
 
 
+def filter_valid_input(
+        input_list: List[torch.Tensor]
+) -> Tuple[List[torch.Tensor], List[bool]]:
+    func_valid = lambda x: x is not None
+    valid_list = list(map(func_valid, input_list))
+    input_list = list(filter(func_valid, input_list))
+    return input_list, valid_list
+
+
+def restore_full_output(output_list: List[torch.Tensor],
+                        valid_list: List[bool]) -> List[torch.Tensor]:
+    index_list = list(accumulate(map(int, valid_list)))
+    output_list = list(
+        map(lambda valid, index: output_list[index - 1]
+            if valid else None, valid_list, index_list))
+    return output_list
+
+
 def allgather(
     input: Union[torch.Tensor, List[torch.Tensor]],
     mapping: Mapping,
@@ -155,8 +174,10 @@ def allgather(
         if isinstance(input, torch.Tensor):
             assert input.shape[dim] == sizes[mapping.tp_rank]
         else:
-            assert all(
-                [val.shape[dim] == sizes[mapping.tp_rank] for val in input])
+            assert all([
+                val.shape[dim] == sizes[mapping.tp_rank] for val in input
+                if val is not None
+            ])
         # 'sizes' is not needed if all inputs in the same TP group have the same shape
         for split_size in sizes[1:]:
             if split_size != sizes[0]:
@@ -170,6 +191,7 @@ def allgather(
         output_info = get_output_info(input, dim)
         input = input.contiguous().view(-1, output_info['numel_base'])
     else:
+        input, valid = filter_valid_input(input)
         torch_op = torch.ops.trtllm.allgather_list
         output_info = [get_output_info(val, dim) for val in input]
         input = [
@@ -202,6 +224,7 @@ def convert_output(x, x_info):
             convert_output(val, val_info)
             for val, val_info in zip(output, output_info)
         ]
+        output = restore_full_output(output, valid)
     return output
 
 
@@ -220,7 +243,10 @@ def reducescatter(
         if isinstance(input, torch.Tensor):
             assert input.shape[dim] == sum_split_size
         else:
-            assert all([val.shape[dim] == sum_split_size for val in input])
+            assert all([
+                val.shape[dim] == sum_split_size for val in input
+                if val is not None
+            ])
         # 'sizes' is not needed if all outputs in the same TP group have the same shape
         for split_size in sizes[1:]:
             if split_size != sizes[0]:
@@ -245,6 +271,7 @@ def convert_input(x, x_info):
         output_info = get_output_info(input, dim)
         input = convert_input(input, output_info)
     else:
+        input, valid = filter_valid_input(input)
         torch_op = torch.ops.trtllm.reducescatter_list
         output_info = [get_output_info(val, dim) for val in input]
         input = [
@@ -265,6 +292,7 @@ def convert_input(x, x_info):
             val.view(val_info['output_shape'])
             for val, val_info in zip(output, output_info)
         ]
+        output = restore_full_output(output, valid)
     return output