Merge branch 'main' into chunked-prefill-spec-dec

Author: mikeiovine

docker/Makefile

@@ -180,7 +180,8 @@ jenkins-aarch64_%: IMAGE_WITH_TAG = $(shell . ../jenkins/current_image_tags.prop
 jenkins-aarch64_%: STAGE = tritondevel
 
 # For x86_64
-jenkins-rockylinux8_%: IMAGE_WITH_TAG = $(shell . ../jenkins/current_image_tags.properties && echo $$LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE)
+jenkins-rockylinux8_%: PYTHON_VERSION_TAG_ID = $(if $(findstring 3.12,${PYTHON_VERSION}),PY312,$(if $(findstring 3.10,${PYTHON_VERSION}),PY310,$(error Unknown PYTHON_VERSION specified)))
+jenkins-rockylinux8_%: IMAGE_WITH_TAG = $(shell . ../jenkins/current_image_tags.properties && echo $$LLM_ROCKYLINUX8_${PYTHON_VERSION_TAG_ID}_DOCKER_IMAGE)
 jenkins-rockylinux8_%: STAGE = tritondevel
 jenkins-rockylinux8_%: BASE_IMAGE = nvidia/cuda
 jenkins-rockylinux8_%: BASE_TAG = 12.9.0-devel-rockylinux8

docker/README.md

@@ -89,13 +89,10 @@ equivalent containers as [described above](#building-docker-images-with-gnu-make
 ### Jenkins Integration
 
 [`Makefile`](Makefile) has special targets for building, pushing and running the Docker build image used on Jenkins.
-The full image name and tag is defined in [`L0_MergeRequest.groovy`](../jenkins/L0_MergeRequest.groovy). The `make`
-system will parse this name as the value of `LLM_DOCKER_IMAGE`. To build and push a new Docker image for Jenkins,
-define a new image name and tag in [`L0_MergeRequest.groovy`](../jenkins/L0_MergeRequest.groovy) and run
+The full image names and tags are defined in [`current_image_tags.properties`](../jenkins/current_image_tags.properties). The `make`
+system will parse the names/tags from this file.
 
-```bash
-make -C docker jenkins_push
-```
+#### Running
 
 Start a new container using the same image as Jenkins using your local user account with
 
@@ -134,6 +131,38 @@ make -C docker trtllm_run LOCAL_USER=1 DOCKER_PULL=1
 The argument `DOCKER_PULL=1` instructs `make` to pull the latest version of the image before deploying it in the container.
 By default, the release images built in the above manner are tagged by their `git` branch name and may be frequently updated.
 
+#### Building CI images
+
+To build and push a new Docker image for Jenkins, define new image names and tags in [`current_image_tags.properties`](../jenkins/current_image_tags.properties) and run
+
+```bash
+# Commands assume an amd64 host
+make -C docker jenkins_build
+#
+docker buildx create --name multi-builder
+make -C docker jenkins-aarch64_build \
+    DOCKER_BUILD_ARGS="--platform arm64 --builder=multi-builder"
+#
+# check jenkins/BuildDockerImage.groovy for current Python versions
+make -C docker jenkins-rockylinux8_build PYTHON_VERSION=3.12.3
+make -C docker jenkins-rockylinux8_build PYTHON_VERSION=3.10.12
+```
+
+The resulting images then need to be pushed:
+
+```bash
+sh -c '. jenkins/current_image_tags.properties && echo $LLM_DOCKER_IMAGE $LLM_SBSA_DOCKER_IMAGE $LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE $LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE' | tr ' ' '\n' | xargs -I{} docker push {}
+```
+
+Alternatively, it is possible to trigger the image build by opening a new pull request and commenting
+
+```text
+/bot run --stage-list "Build-Docker-Images"
+```
+
+The resulting images can then be re-tagged using `scripts/rename_docker_images.py`
+and the new tags included in [`current_image_tags.properties`](../jenkins/current_image_tags.properties).
+
 ### Docker rootless
 
 Some aspects require special treatment when using [Docker rootless mode](https://docs.docker.com/engine/security/rootless/). The `docker/Makefile` contains heuristics to detect Docker rootless mode. When assuming

jenkins/current_image_tags.properties

@@ -8,7 +8,10 @@
 # NB: Although string interpolation is supported, redundant substrings are
 #     kept in the variables below for interoperability with
 #     scripts/rename_docker_images.py
-LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.05-py3-x86_64-ubuntu24.04-trt10.11.0.33-skip-tritondevel-202507150652-9504
-LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.05-py3-aarch64-ubuntu24.04-trt10.11.0.33-skip-tritondevel-202507150652-9504
-LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py310-trt10.11.0.33-skip-tritondevel-202507150652-9504
-LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py312-trt10.11.0.33-skip-tritondevel-202507150652-9504
+#
+# NB: Typically, the suffix indicates the PR whose CI pipeline generated the images. In case that
+#     images are adopted from PostMerge pipelines, the abbreviated commit hash is used instead.
+LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.05-py3-x86_64-ubuntu24.04-trt10.11.0.33-skip-tritondevel-202507162011-ec3ebae
+LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.05-py3-aarch64-ubuntu24.04-trt10.11.0.33-skip-tritondevel-202507162011-ec3ebae
+LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py310-trt10.11.0.33-skip-tritondevel-202507162011-ec3ebae
+LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py312-trt10.11.0.33-skip-tritondevel-202507162011-ec3ebae

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -966,19 +966,14 @@ def _executor_loop(self):
                         self._prepare_disagg_gen_transmission_complete(
                             scheduled_batch)
 
+                        # Return the first token to the client
+                        self._handle_first_token_response(scheduled_batch)
+
                     self.resource_manager.prepare_resources(scheduled_batch)
                     if self.drafter is not None:
                         self.drafter.prepare_draft_tokens(
                             scheduled_batch, self.resource_manager)
 
-                    if self.kv_cache_transceiver:
-                        # For generation requests which have completed KV cache transfer
-                        self._prepare_disagg_gen_transmission_complete(
-                            scheduled_batch)
-
-                        # Return the first token to the client
-                        self._handle_first_token_response(scheduled_batch)
-
                     batch_outputs = self._forward_step(scheduled_batch)
 
                     if self.guided_decoder is not None: