From 678de449a1d2f36c9b5059e2bc39adf0275782b0 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:24:14 -0700 Subject: [PATCH 1/9] [NVBUG-6448152][test] apply audited current-runtime overlay Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- docker/common/install_nixl.sh | 8 ++++---- jenkins/current_image_tags.properties | 10 +++++----- requirements-dev.txt | 4 ++-- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/docker/common/install_nixl.sh b/docker/common/install_nixl.sh index f14398ced80c..51208e0ea89d 100644 --- a/docker/common/install_nixl.sh +++ b/docker/common/install_nixl.sh @@ -4,7 +4,7 @@ set -ex GITHUB_URL="https://github.com" UCX_INSTALL_PATH="/usr/local/ucx/" CUDA_PATH="/usr/local/cuda" -NIXL_VERSION="v1.0.1" +NIXL_VERSION="v1.3.1" NIXL_REPO="https://github.com/ai-dynamo/nixl.git" OLD_LD_LIBRARY_PATH=$LD_LIBRARY_PATH @@ -23,9 +23,6 @@ pip3 install meson ninja pybind11 setuptools git clone --depth 1 -b ${NIXL_VERSION} ${NIXL_REPO} cd nixl -# Remove POSIX backend compilation from meson.build -sed -i "/^subdir('posix')/d" src/plugins/meson.build - CUDA_SO_PATH=$(find "/usr/local" -name "libcuda.so.1" 2>/dev/null | head -n1) if [[ -z "$CUDA_SO_PATH" ]]; then @@ -42,6 +39,9 @@ meson setup builddir \ -Dcudapath_inc="$CUDA_PATH/include" \ -Dgds_path="$GDS_PATH" \ -Dinstall_headers=true \ + -Ddisable_plugins=POSIX \ + -Dbuild_tests=false \ + -Dbuild_examples=false \ --buildtype=release cd builddir && ninja install diff --git a/jenkins/current_image_tags.properties b/jenkins/current_image_tags.properties index c33159d0bfe5..504fd8f234a8 100644 --- a/jenkins/current_image_tags.properties +++ b/jenkins/current_image_tags.properties @@ -13,8 +13,8 @@ # images are adopted from PostMerge pipelines, the abbreviated commit hash is used instead. IMAGE_NAME=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm -LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.02-py3-x86_64-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202606091844-14706 -LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.02-py3-sbsa-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202606091844-14706 -LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-rocky8-x86_64-rocky8-py310-trt10.15.1.29-skip-tritondevel-202606091844-14706 -LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-rocky8-x86_64-rocky8-py312-trt10.15.1.29-skip-tritondevel-202606091844-14706 -LLM_SBSA_WHEEL_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-trt10.15.1.29-skip-tritondevel-202606091844-14706 +LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.05-py3-x86_64-ubuntu24.04-trt10.16.1.11-skip-tritondevel-202607151440-16194 +LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.05-py3-sbsa-ubuntu24.04-trt10.16.1.11-skip-tritondevel-202607151440-16194 +LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.2.1-devel-rocky8-x86_64-rocky8-py310-trt10.16.1.11-skip-tritondevel-202607151440-16194 +LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.2.1-devel-rocky8-x86_64-rocky8-py312-trt10.16.1.11-skip-tritondevel-202607151440-16194 +LLM_SBSA_WHEEL_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.2.1-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-trt10.16.1.11-skip-tritondevel-202607151440-16194 diff --git a/requirements-dev.txt b/requirements-dev.txt index 554a1bf58d46..4cef09afda7f 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -39,9 +39,9 @@ opentelemetry-api>=1.26.0 opentelemetry-exporter-otlp>=1.26.0 opentelemetry-semantic-conventions-ai>=0.4.1 fuzzywuzzy==0.18.0 -aiperf==0.6.0 +aiperf==0.8.0 nanobind>=2.9.0 -nixl-cu13==1.0.1 +nixl-cu13==1.3.1 cupti-python>=13.0,<13.2 nvidia-cuda-cupti>=13.0,<13.2 cxxfilt From aea9d7618c7959c47dea7b5008c40966dd5982d1 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:24:49 -0700 Subject: [PATCH 2/9] [NVBUG-6448152][test] hold current CI harness constant Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_Test.groovy | 123 +++++++++++++++++------- jenkins/scripts/perf/local/slurm_run.sh | 9 ++ jenkins/scripts/slurm_run.sh | 9 ++ 3 files changed, 104 insertions(+), 37 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index f05f476b6b9f..b7b78fc5971c 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -1,3 +1,19 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + @Library(['bloom-jenkins-shared-lib@main', 'trtllm-jenkins-shared-lib@main']) _ import java.lang.InterruptedException @@ -46,7 +62,7 @@ LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE = env.wheelDockerImagePy312 LLM_WHEEL_DOCKER_IMAGE = env.wheelDockerImage // DLFW torch image -DLFW_IMAGE = "urm.nvidia.com/docker/nvidia/pytorch:26.02-py3" +DLFW_IMAGE = "urm.nvidia.com/docker/nvidia/pytorch:26.04-py3" //Ubuntu base image UBUNTU_22_04_IMAGE = "urm.nvidia.com/docker/ubuntu:22.04" @@ -925,14 +941,18 @@ def runLLMTestlistWithAgent(pipeline, platform, testList, config=VANILLA_CONFIG, throw e } } finally { - captureSlurmJobNodeList(pipeline, cluster, partition.clusterName, slurmJobID, placementContext, stageName) - stage("Clean Up Slurm Resource") { - // Workaround to handle the interruption during clean up SLURM resources - retry(3) { - try { - cleanUpNodeResources(pipeline, cluster, partition.clusterName, nodeName, slurmJobID) - } catch (Exception e) { - error "Error during clean up SLURM resources: ${e.getMessage()} and retrying." + // Resource cleanup must run even if SLURM metadata capture is interrupted. + try { + captureSlurmJobNodeList(pipeline, cluster, partition.clusterName, slurmJobID, placementContext, stageName) + } finally { + stage("Clean Up Slurm Resource") { + // Workaround to handle the interruption during clean up SLURM resources + retry(3) { + try { + cleanUpNodeResources(pipeline, cluster, partition.clusterName, nodeName, slurmJobID) + } catch (Exception e) { + error "Error during clean up SLURM resources: ${e.getMessage()} and retrying." + } } } } @@ -1100,6 +1120,8 @@ def getMountListForSlurmTest(SlurmCluster cluster, boolean useSbatch = false) throw new Exception("Unsupported container runtime: ${cluster.containerRuntime}") } + // TODO: Add mounts for different cache directories like pip, triton, etc. + return mounts } @@ -1372,6 +1394,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG "--container-image=$containerImageArg", "--container-workdir=$jobWorkspace", "--container-mounts=$mounts", + "--no-container-mount-home", "--container-env=NVIDIA_IMEX_CHANNELS" ] envVarsToExport.each { varName, varValue -> @@ -1463,7 +1486,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG """ } else { if(nodeCount > 1) { - srunArgs.add("--mpi=pmi2") + srunArgs.add("--mpi=pmix") } def scriptContent = """ @@ -1688,15 +1711,19 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG stageIsInterrupted = true throw e } finally { - captureSlurmJobNodeList(pipeline, cluster, partition.clusterName, placementContext?.slurmJobId ?: null, placementContext, stageName, jobWorkspace) - uploadResults(pipeline, cluster, partition.clusterName, jobUID, stageName, stageIsInterrupted, postTag) - stage("Clean Up Slurm Resource") { - // Workaround to handle the interruption during clean up SLURM resources - retry(3) { - try { - cleanUpSlurmResources(pipeline, cluster, partition.clusterName, jobUID) - } catch (Exception e) { - error "Error during clean up SLURM resources: ${e.getMessage()} and retrying." + // Resource cleanup must run even if metadata capture or result upload is interrupted. + try { + captureSlurmJobNodeList(pipeline, cluster, partition.clusterName, placementContext?.slurmJobId ?: null, placementContext, stageName, jobWorkspace) + uploadResults(pipeline, cluster, partition.clusterName, jobUID, stageName, stageIsInterrupted, postTag) + } finally { + stage("Clean Up Slurm Resource") { + // Workaround to handle the interruption during clean up SLURM resources + retry(3) { + try { + cleanUpSlurmResources(pipeline, cluster, partition.clusterName, jobUID) + } catch (Exception e) { + error "Error during clean up SLURM resources: ${e.getMessage()} and retrying." + } } } } @@ -2262,7 +2289,7 @@ def cacheErrorAndUploadResult(stageName, taskRunner, finallyRunner, noResultIfSu if (stageIsFailed && !suppressTestReporting) { if (stageIsInterrupted) { echo "Stage is interrupted, skip to generate terminated unexpectedly test result." - } else { + } else if (!fileExists("${stageName}/results-timeout.xml")) { // Generate timeout test result xml if there are terminated unexpectedly tests generateTimeoutTestResultXml(pipeline, stageName) } @@ -2474,6 +2501,12 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod - SYS_ADMIN""" break } + // Temporarily avoid an arm64 CPU builder with repeated pod DNS/JNLP failures seen in Build-SBSA #5564. + def blockedNodeAffinity = targetCloud == "kubernetes-cpu" && arch == "arm64" ? ''' + - key: "kubernetes.io/hostname" + operator: NotIn + values: + - "rl300-0021.ipp2a1.colossus"''' : "" def nodeLabel = trtllm_utils.generateNodeLabel(nodeLabelPrefix) def pvcVolume = """ - name: sw-tensorrt-pvc @@ -2529,6 +2562,7 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod values: - "core" - "qa_only" +${blockedNodeAffinity} nodeSelector: ${selectors} containers: ${containerConfig} @@ -2652,7 +2686,7 @@ def launchTestListCheck(pipeline) def llmPath = sh (script: "realpath .", returnStdout: true).trim() def llmSrc = "${llmPath}/TensorRT-LLM/src" trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install -r ${llmSrc}/requirements-dev.txt") - sh "NVIDIA_TRITON_SERVER_VERSION=26.02 LLM_ROOT=${llmSrc} LLM_BACKEND_ROOT=${llmSrc}/triton_backend python3 ${llmSrc}/scripts/check_test_list.py --l0 --qa --waive" + sh "NVIDIA_TRITON_SERVER_VERSION=26.04 LLM_ROOT=${llmSrc} LLM_BACKEND_ROOT=${llmSrc}/triton_backend python3 ${llmSrc}/scripts/check_test_list.py --l0 --qa --waive" } catch (InterruptedException e) { throw e } catch (Exception e) { @@ -3483,7 +3517,7 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO } trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmSrc} && pip3 install -r requirements-dev.txt") if (stageName.contains("-Ray-")) { - trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install ray[default]==2.54.1") + trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install ray[default]==2.55.1") } if (!skipInstallWheel) { trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmPath} && pip3 install --force-reinstall --no-deps TensorRT-LLM/tensorrt_llm-*.whl") @@ -3694,6 +3728,13 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO error "Regular tests failed after rerun attempt" } rerunFailed = true + } else if (generateTimeoutTestResultXml(pipeline, stageName)) { + // Rerun passed but the first run had a timeout: mark this + // stage FAILURE so "[${stageName}] Run Pytest" turns red, + // not just the enclosing parent stage. + catchError(buildResult: 'SUCCESS', stageResult: 'FAILURE') { + error "Some tests terminated unexpectedly, please check the test report." + } } } @@ -3725,6 +3766,10 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO error "Some tests still failed after rerun attempts, please check the test report." } + if (fileExists("${stageName}/results-timeout.xml") || generateTimeoutTestResultXml(pipeline, stageName)) { + error "Some tests terminated unexpectedly, please check the test report." + } + if (perfMode) { // Only PyTorch perf stages remain; the TensorRT perf baseline was removed. basePerfFilename = "base_perf_pytorch.csv" @@ -4397,6 +4442,7 @@ def launchTestJobs(pipeline, testFilter) // "L40S-TensorRT-Post-Merge-4": ["l40s", "l0_l40s", 4, 5], // "L40S-TensorRT-Post-Merge-5": ["l40s", "l0_l40s", 5, 5], "L40S-FMHA-Post-Merge-1": ["l40s", "l0_l40s", 1, 1], + "H100_PCIe-AutoDeploy-Post-Merge-1": ["h100-cr", "l0_h100", 1, 1], "H100_PCIe-CPP-Post-Merge-1": ["h100-cr", "l0_h100", 1, 1], // "H100_PCIe-TensorRT-Post-Merge-1": ["h100-cr", "l0_h100", 1, 5], // "H100_PCIe-TensorRT-Post-Merge-2": ["h100-cr", "l0_h100", 2, 5], @@ -4407,7 +4453,6 @@ def launchTestJobs(pipeline, testFilter) // "B200_PCIe-TensorRT-Post-Merge-1": ["b100-ts2", "l0_b200", 1, 2], // "B200_PCIe-TensorRT-Post-Merge-2": ["b100-ts2", "l0_b200", 2, 2], "H100_PCIe-PyTorch-Perf-1": ["h100-cr", "l0_perf", 1, 1], - "DGX_H200-4_GPUs-Triton-Post-Merge-1": ["dgx-h200-x4", "l0_dgx_h200", 1, 1, 4], "DGX_H200-8_GPUs-PyTorch-Post-Merge-1": ["dgx-h200-x8", "l0_dgx_h200", 1, 1, 8], "DGX_H200-4_GPUs-PyTorch-Post-Merge-1": ["dgx-h200-x4", "l0_dgx_h200", 1, 1, 4], "DGX_H200-8_GPUs-PyTorch-PerfSanity-Post-Merge-1": ["dgx-h200-x8", "l0_dgx_h200_perf_sanity", 1, 1, 8], @@ -4447,8 +4492,6 @@ def launchTestJobs(pipeline, testFilter) "DGX_H100-PyTorch-6": ["auto:dgx-h100-x1", "l0_h100", 6, 6], "DGX_H100-PyTorch-Post-Merge-1": ["auto:dgx-h100-x1", "l0_h100", 1, 2], "DGX_H100-PyTorch-Post-Merge-2": ["auto:dgx-h100-x1", "l0_h100", 2, 2], - "DGX_A100-Triton-Post-Merge-1": ["auto:dgx-a100-x1", "l0_a100", 1, 2], - "DGX_A100-Triton-Post-Merge-2": ["auto:dgx-a100-x1", "l0_a100", 2, 2], "DGX_A100-FMHA-Post-Merge-1": ["auto:dgx-a100-x1", "l0_a100", 1, 1], "DGX_H100-2_GPUs-PyTorch-Others-1": ["auto:dgx-h100-x2", "l0_dgx_h100", 1, 2, 2], "DGX_H100-2_GPUs-PyTorch-Others-2": ["auto:dgx-h100-x2", "l0_dgx_h100", 2, 2, 2], @@ -4461,6 +4504,7 @@ def launchTestJobs(pipeline, testFilter) "DGX_H100-4_GPUs-PyTorch-Ray-1": ["auto:dgx-h100-x4", "l0_dgx_h100", 1, 1, 4], "DGX_H100-4_GPUs-AutoDeploy-1": ["auto:dgx-h100-x4", "l0_dgx_h100", 1, 1, 4], "DGX_H100-4_GPUs-AutoDeploy-Post-Merge-1": ["auto:dgx-h100-x4", "l0_dgx_h100", 1, 1, 4], + "DGX_H100-4_GPUs-PyTorch-Post-Merge-1": ["auto:dgx-h100-x4", "l0_dgx_h100", 1, 1, 4], "DGX_B200-PyTorch-1": ["auto:dgx-b200-flex", "l0_b200", 1, 9, 1, 1, true], "DGX_B200-PyTorch-2": ["auto:dgx-b200-flex", "l0_b200", 2, 9, 1, 1, true], "DGX_B200-PyTorch-3": ["auto:dgx-b200-flex", "l0_b200", 3, 9, 1, 1, true], @@ -4471,7 +4515,7 @@ def launchTestJobs(pipeline, testFilter) "DGX_B200-PyTorch-8": ["auto:dgx-b200-flex", "l0_b200", 8, 9, 1, 1, true], "DGX_B200-PyTorch-9": ["auto:dgx-b200-flex", "l0_b200", 9, 9, 1, 1, true], "DGX_B200-AutoDeploy-1": ["auto:dgx-b200-flex", "l0_b200", 1, 1, 1, 1, true], - "DGX_B200-Triton-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200", 1, 1, 1, 1, true], + "DGX_B200-AutoDeploy-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200", 1, 1, 1, 1, true], "DGX_B200-PyTorch-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200", 1, 2, 1, 1, true], "DGX_B200-PyTorch-Post-Merge-2": ["auto:dgx-b200-flex", "l0_b200", 2, 2, 1, 1, true], "DGX_B200-2_GPUs-PyTorch-1": ["auto:dgx-b200-flex", "l0_dgx_b200", 1, 1, 2, 1, true], @@ -4480,17 +4524,22 @@ def launchTestJobs(pipeline, testFilter) "DGX_B200-4_GPUs-PyTorch-3": ["auto:dgx-b200-flex", "l0_dgx_b200", 3, 3, 4, 1, true], "DGX_B200-4_GPUs-PyTorch-Ray-1": ["auto:dgx-b200-flex", "l0_dgx_b200", 1, 1, 4, 1, true], "DGX_B200-4_GPUs-AutoDeploy-1": ["auto:dgx-b200-flex", "l0_dgx_b200", 1, 1, 4, 1, true], + "DGX_B200-4_GPUs-AutoDeploy-Post-Merge-1": ["auto:dgx-b200-flex", "l0_dgx_b200", 1, 1, 4, 1, true], "DGX_B200-4_GPUs-PyTorch-Post-Merge-1": ["auto:dgx-b200-flex", "l0_dgx_b200", 1, 4, 4, 1, true], "DGX_B200-4_GPUs-PyTorch-Post-Merge-2": ["auto:dgx-b200-flex", "l0_dgx_b200", 2, 4, 4, 1, true], "DGX_B200-4_GPUs-PyTorch-Post-Merge-3": ["auto:dgx-b200-flex", "l0_dgx_b200", 3, 4, 4, 1, true], "DGX_B200-4_GPUs-PyTorch-Post-Merge-4": ["auto:dgx-b200-flex", "l0_dgx_b200", 4, 4, 4, 1, true], - "DGX_B200-8_GPUs-PyTorch-1": ["auto:dgx-b200-flex", "l0_dgx_b200", 1, 3, 8, 1, true], - "DGX_B200-8_GPUs-PyTorch-2": ["auto:dgx-b200-flex", "l0_dgx_b200", 2, 3, 8, 1, true], - "DGX_B200-8_GPUs-PyTorch-3": ["auto:dgx-b200-flex", "l0_dgx_b200", 3, 3, 8, 1, true], + "DGX_B200-8_GPUs-PyTorch-1": ["auto:dgx-b200-flex", "l0_dgx_b200", 1, 4, 8, 1, true], + "DGX_B200-8_GPUs-PyTorch-2": ["auto:dgx-b200-flex", "l0_dgx_b200", 2, 4, 8, 1, true], + "DGX_B200-8_GPUs-PyTorch-3": ["auto:dgx-b200-flex", "l0_dgx_b200", 3, 4, 8, 1, true], + "DGX_B200-8_GPUs-PyTorch-4": ["auto:dgx-b200-flex", "l0_dgx_b200", 4, 4, 8, 1, true], "DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1": ["auto:dgx-b200-flex", "l0_dgx_b200", 1, 1, 8, 1, true], - "DGX_B200-4_GPUs-Verl-Post-Merge-1": ["auto:dgx-b200-flex", "l0_verl", 1, 1, 4, 1, true], + // Disable Verl stage due to https://nvbugs/6236818. + // Please re-enable it after the bug is fixed. + // "DGX_B200-4_GPUs-Verl-Post-Merge-1": ["auto:dgx-b200-flex", "l0_verl", 1, 1, 4, 1, true], "B300-PyTorch-1": ["auto:dgx-b300-flex", "l0_b300", 1, 2, 1, 1, true], "B300-PyTorch-2": ["auto:dgx-b300-flex", "l0_b300", 2, 2, 1, 1, true], + "B300-PyTorch-Post-Merge-1": ["auto:dgx-b300-flex", "l0_b300", 1, 1, 1, 1, true], "DGX_B300-4_GPUs-PyTorch-1": ["auto:dgx-b300-flex", "l0_dgx_b300", 1, 1, 4, 1, true], "DGX_B300-4_GPUs-PyTorch-Post-Merge-1": ["auto:dgx-b300-flex", "l0_dgx_b300", 1, 2, 4, 1, true], "DGX_B300-4_GPUs-PyTorch-Post-Merge-2": ["auto:dgx-b300-flex", "l0_dgx_b300", 2, 2, 4, 1, true], @@ -4517,7 +4566,7 @@ def launchTestJobs(pipeline, testFilter) x86SlurmTestConfigs = cbtsResizeSplits(x86SlurmTestConfigs) fullSet += x86SlurmTestConfigs.keySet() - parallelSlurmJobs = x86SlurmTestConfigs.collectEntries{key, values -> [key, [createKubernetesPodConfig(LLM_DOCKER_IMAGE, "slurm", "amd64"), { attemptTag, isFinalAttempt, retryContext = null -> + parallelSlurmJobs = x86SlurmTestConfigs.collectEntries{key, values -> [key, [createKubernetesPodConfig(LLM_DOCKER_IMAGE.replace("aarch64", "x86_64"), "slurm", "amd64"), { attemptTag, isFinalAttempt, retryContext = null -> // attemptTag comes from runKubernetesPodWithInfraRetry for the outer // dispatcher pod (when retry is enabled — see opts below) and is // threaded into runLLMTestlistOnSlurm so a future re-enable of outer @@ -4789,7 +4838,7 @@ def launchTestJobs(pipeline, testFilter) // singleAttempt:true disables the outer K8s pod retry; see the x86 // SLURM closure above for the full rationale (cap nested retry budget // so consistently-timing-out tests don't burn ~36h on retry cascades). - parallelSlurmJobs = SBSASlurmTestConfigs.collectEntries{key, values -> [key, [createKubernetesPodConfig(LLM_DOCKER_IMAGE, "slurm", "arm64"), { attemptTag, isFinalAttempt, retryContext = null -> + parallelSlurmJobs = SBSASlurmTestConfigs.collectEntries{key, values -> [key, [createKubernetesPodConfig(LLM_DOCKER_IMAGE.replace("aarch64", "x86_64"), "slurm", "amd64"), { attemptTag, isFinalAttempt, retryContext = null -> // attemptTag is threaded into runLLMTestlistOnSlurm as the outer // dispatcher pod's tag so the inner SLURM retry's postTag can't // collide with a previous dispatcher pod's upload. See the x86 @@ -4807,7 +4856,7 @@ def launchTestJobs(pipeline, testFilter) // Add SBSA multi node Slurm jobs // singleAttempt:true disables the outer K8s pod retry; see above. - parallelMultiNodesSBSAJobs = multiNodesSBSAConfigs.collectEntries{key, values -> [key, [createKubernetesPodConfig(LLM_DOCKER_IMAGE, "slurm", "arm64"), { attemptTag, isFinalAttempt, retryContext = null -> + parallelMultiNodesSBSAJobs = multiNodesSBSAConfigs.collectEntries{key, values -> [key, [createKubernetesPodConfig(LLM_DOCKER_IMAGE.replace("aarch64", "x86_64"), "slurm", "amd64"), { attemptTag, isFinalAttempt, retryContext = null -> def config = LINUX_AARCH64_CONFIG if (key.contains("single-device")) { config = SINGLE_DEVICE_CONFIG @@ -4996,14 +5045,14 @@ def launchTestJobs(pipeline, testFilter) def platform = cpu_arch == X86_64_TRIPLE ? "x86_64" : "sbsa" trtllm_utils.llmExecStepWithRetry(pipeline, script: "wget https://developer.download.nvidia.com/compute/cuda/repos/${ubuntu_version}/${platform}/cuda-keyring_1.1-1_all.deb") trtllm_utils.llmExecStepWithRetry(pipeline, script: "dpkg -i cuda-keyring_1.1-1_all.deb") - trtllm_utils.llmExecStepWithRetry(pipeline, script: "apt-get update && apt-get install -y cuda-toolkit-13-1") + trtllm_utils.llmExecStepWithRetry(pipeline, script: "apt-get update && apt-get install -y cuda-toolkit-13-2") } - // Extra PyTorch CUDA 13.0 install for all bare-metal environments (Default PyTorch is for CUDA 12.8) + // Extra PyTorch CUDA 13.2 install for all bare-metal environments (Default PyTorch is for CUDA 12.8) if (values[6]) { - echo "###### Extra PyTorch CUDA 13.0 install Start ######" + echo "###### Extra PyTorch CUDA 13.2 install Start ######" // Use internal mirror instead of https://download.pytorch.org/whl/cu130 for better network stability. // PyTorch CUDA 13.0 package and torchvision package can be installed as expected. - trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install torch==2.10.0+cu130 torchvision==0.25.0+cu130 --extra-index-url https://urm.nvidia.com/artifactory/api/pypi/pytorch-cu128-remote/simple --extra-index-url https://download.pytorch.org/whl/cu130") + trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install torch==2.11.0+cu130 torchvision==0.26.0+cu130 --extra-index-url https://urm.nvidia.com/artifactory/api/pypi/pytorch-cu128-remote/simple --extra-index-url https://download.pytorch.org/whl/cu130") } def libEnv = [] diff --git a/jenkins/scripts/perf/local/slurm_run.sh b/jenkins/scripts/perf/local/slurm_run.sh index d10dc6fc23a8..a874b4fb965d 100755 --- a/jenkins/scripts/perf/local/slurm_run.sh +++ b/jenkins/scripts/perf/local/slurm_run.sh @@ -15,6 +15,15 @@ fi cd $llmSrcNode/tests/integration/defs +# Force PMIx to use the in-memory hash GDS instead of ds12/ds21 shared-memory. +# Under `srun --mpi=pmix` with the DLFW 26.04 OpenMPI build, the shared-memory +# GDS modes can fail to publish UCX worker addresses across nodes, producing: +# pml_ucx.c:178 Error: Failed to receive UCX worker address: Not found (-13) +# pml_ucx.c:482 Error: Failed to resolve UCX endpoint for rank N +# See https://github.com/open-mpi/ompi/issues/6981. Setting this is a no-op +# when PMIx isn't used. +export PMIX_MCA_gds=hash + # Turn off "exit on error" so the following lines always run set +e diff --git a/jenkins/scripts/slurm_run.sh b/jenkins/scripts/slurm_run.sh index 686bb80f1d16..1b9391d20fca 100755 --- a/jenkins/scripts/slurm_run.sh +++ b/jenkins/scripts/slurm_run.sh @@ -88,6 +88,15 @@ if [ "${UCX_TLS:-}" = "tcp" ]; then unset UCX_TLS echo "Unset UCX_TLS (cluster injected UCX_TLS=tcp)" fi + +# Force PMIx to use the in-memory hash GDS instead of ds12/ds21 shared-memory. +# Under `srun --mpi=pmix` with the DLFW 26.04 OpenMPI build, the shared-memory +# GDS modes can fail to publish UCX worker addresses across nodes, producing: +# pml_ucx.c:178 Error: Failed to receive UCX worker address: Not found (-13) +# pml_ucx.c:482 Error: Failed to resolve UCX endpoint for rank N +# See https://github.com/open-mpi/ompi/issues/6981. Setting this is a no-op +# when PMIx isn't used. +export PMIX_MCA_gds=hash echo "Library Path:" echo "$LD_LIBRARY_PATH" env | sort From 1eefb08b703007041a0d02b1f0c7648edd3924c0 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Mon, 6 Jul 2026 13:50:37 +0800 Subject: [PATCH 3/9] [None][fix] Fix marlin_nvfp4_template.h compilation error (#15660) Signed-off-by: Yihan Wang --- .../kernels/marlin/marlin_nvfp4_template.h | 48 +++++++------------ 1 file changed, 17 insertions(+), 31 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h index fa02e8387430..9fdc9d2a6d38 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -52,40 +52,26 @@ namespace MARLIN_NAMESPACE_NAME // Empty kernel stub for non-Hopper device passes; see marlin.cuh. #if defined(__CUDA_ARCH__) && !(__CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000) -template shared - // fetch pipeline - bool const has_act_order, // whether act_order is enabled - int group_blocks // number of consecutive 16x16 blocks - // with a separate quantization scale - // (implicit const: trailing NTTP) +template shared + // fetch pipeline + int group_blocks // number of consecutive 16x16 blocks + // with a separate quantization scale + // (implicit const: trailing NTTP) > -__global__ void Marlin(int4 const* __restrict__ A, // fp16 input matrix of shape mxk - int4 const* __restrict__ B, // 4bit quantized weight matrix of shape kxn - int4* __restrict__ C, // fp16 output buffer of shape mxn - int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) - int4 const* __restrict__ scales_ptr, // fp16 quantization scales of shape - // (k/groupsize)xn - int const* __restrict__ g_idx, // int32 group indices of shape k - int num_groups, // number of scale groups per output channel - int prob_m, // batch dimension m - int prob_n, // output dimension n - int prob_k, // reduction dimension k - int* locks, // extra global storage for barrier synchronization - bool use_fp32_reduce // whether to use fp32 global reduce -) +__global__ void Marlin(MARLIN_KERNEL_PARAMS) { } -} // namespace marlin +} // namespace MARLIN_NAMESPACE_NAME #else From d2f5d4ba3fd5c2c046f99fcde6e5cb4f23082ffe Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 24 Jul 2026 09:40:31 -0700 Subject: [PATCH 4/9] [NVBUG-6448152][test] normalize CMake context only Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> (cherry picked from commit 6fc7f330fcd621ddc632a56614f58f3208707cba) --- cpp/tensorrt_llm/batch_manager/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index 88e2484cc6f1..128a24ee21e1 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -32,11 +32,11 @@ set(SRCS contextProgress.cpp dataTransceiver.cpp decoderBuffers.cpp + kvCacheManager.cpp encoderBuffers.cpp guidedDecoder.cpp handleContextLogits.cpp handleGenerationLogits.cpp - kvCacheManager.cpp kvCacheEventManager.cpp kvCacheTransferManager.cpp kvCacheManagerV2Utils.cpp From b1576daa300ee6dabe232c44202be3610f337dcb Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 24 Jul 2026 09:40:57 -0700 Subject: [PATCH 5/9] [NVBUG-6448152][test] apply semantically matched pre-cancellation consensus factor Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> (cherry picked from commit 9c8bf5952cc2d04d69efe309e42f010d74dc7053) --- .../batch_manager/cacheTransceiver.h | 22 + .../contextTransferCoordinator.h | 109 ++++ .../tensorrt_llm/runtime/utils/mpiTags.h | 8 +- .../tensorrt_llm/runtime/utils/mpiUtils.h | 13 +- cpp/tensorrt_llm/batch_manager/CMakeLists.txt | 3 +- .../batch_manager/cacheTransceiver.cpp | 103 +++- .../contextTransferCoordinator.cpp | 475 ++++++++++++++++++ .../unit_tests/batch_manager/CMakeLists.txt | 1 + .../contextTransferCoordinatorTest.cpp | 143 ++++++ .../multi_gpu/cacheTransceiverTest.cpp | 92 ++++ .../unit_tests/multi_gpu/mpiUtilsTest.cpp | 27 +- 11 files changed, 979 insertions(+), 17 deletions(-) create mode 100644 cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h create mode 100644 cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp create mode 100644 cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 80a06ea5ddf7..30fe338baa52 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -26,6 +26,7 @@ #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" +#include #include #include #include @@ -54,6 +55,7 @@ class BaseKVCacheManager; class CacheSender; class CacheReceiver; +class ContextTransferCoordinator; class CacheTransceiverComm { @@ -148,6 +150,25 @@ class CacheTransceiverComm TLLM_THROW("Input arguments only supported in mpi"); } + [[nodiscard]] std::unique_ptr sendAsync( + void const* buffer, std::size_t size, mpi::MpiType dtype, int dest, mpi::MpiTag tag) const + { + TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI."); + return mMpiComm->sendAsync(buffer, size, dtype, dest, tag); + } + + [[nodiscard]] bool iprobe(int source, mpi::MpiTag tag, MPI_Status* status) const + { + TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI."); + return mMpiComm->iprobe(source, tag, status); + } + + void recv(void* buffer, std::size_t size, mpi::MpiType dtype, int source, mpi::MpiTag tag) const + { + TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI."); + static_cast(mMpiComm->recv(buffer, size, dtype, source, tag)); + } + CacheTransceiverComm split(int color, int key) { if (isMpi()) @@ -298,6 +319,7 @@ class CacheTransceiver : public BaseCacheTransceiver std::shared_ptr mGroupComm; std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; + std::unique_ptr mContextTransferCoordinator; executor::kv_cache::CommState const* mCommState; std::unique_ptr mCacheState; diff --git a/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h b/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h new file mode 100644 index 000000000000..9b156c8887a0 --- /dev/null +++ b/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h @@ -0,0 +1,109 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager +{ + +class CacheTransceiverComm; + +enum class ContextTransferVote : std::uint64_t +{ + kCompleted = 1, + kFailed = 2, +}; + +struct ContextTransferConsensusResult +{ + std::unordered_set completedRequestIds; + std::unordered_set failedRequestIds; + std::unordered_set timedOutRequestIds; +}; + +//! Accumulates one immutable terminal vote per participant and request. +class ContextTransferVoteReducer +{ +public: + explicit ContextTransferVoteReducer(int participantCount); + + void recordVote(int participantRank, std::uint64_t requestId, ContextTransferVote vote); + + //! Record a sticky nonterminal timeout proposal. The final outcome still waits for every terminal vote. + void recordTimeout(std::uint64_t requestId); + + [[nodiscard]] ContextTransferConsensusResult takeReady(); + + void clear() noexcept; + +private: + struct RequestVotes + { + explicit RequestVotes(int participantCount) + : votes(static_cast(participantCount), 0) + { + } + + std::vector votes; + int terminalCount{0}; + bool failed{false}; + bool timedOut{false}; + bool timeoutPending{false}; + }; + + int mParticipantCount; + std::unordered_map mRequestVotes; +}; + +//! Coordinates asynchronous context-transfer outcomes across a rank group. +//! +//! Timeout and terminal events share one ordered stream toward the coordinator. Timeout updates and final commits +//! share another ordered stream toward followers, so a request that times out can never commit success first. +class ContextTransferCoordinator +{ +public: + explicit ContextTransferCoordinator(std::shared_ptr comm); + ~ContextTransferCoordinator(); + + ContextTransferCoordinator(ContextTransferCoordinator const&) = delete; + ContextTransferCoordinator& operator=(ContextTransferCoordinator const&) = delete; + + //! Publish this rank's immutable local terminal outcome without waiting for peers. + void publishLocalOutcome(std::uint64_t requestId, bool failed); + + //! Publish an idempotent, sticky, nonterminal timeout proposal without waiting for peers. + void publishTimeout(std::uint64_t requestId); + + //! Make nonblocking protocol progress and return newly committed global outcomes. + [[nodiscard]] ContextTransferConsensusResult poll(); + + //! Exchange ordered close markers so no active MPI request outlives its backing buffer. Failure aborts closed. + void shutdown() noexcept; + +private: + class Impl; + std::unique_ptr mImpl; +}; + +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h index 32c086c84ee9..49c50e49b2e2 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -71,7 +71,11 @@ enum class MpiTag : int // KvCacheEventManager kKvCacheEventSize = 1026, - kKvCacheEvent = 1027 + kKvCacheEvent = 1027, + + // Asynchronous context-transfer coordination. + kContextTransferEvent = 1028, + kContextTransferUpdate = 1029 }; } // namespace tensorrt_llm::mpi diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h b/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h index 75ec7a534815..b474b4b12cbf 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -235,6 +235,17 @@ class MpiRequest #endif } + [[nodiscard]] bool isCompleted() + { +#if ENABLE_MULTI_DEVICE + int completed = 0; + TLLM_MPI_CHECK(MPI_Test(&mRequest, &completed, MPI_STATUS_IGNORE)); + return completed != 0; +#else + TLLM_THROW("Multi device support is disabled."); +#endif + } + void cancel() { #if ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index 128a24ee21e1..154b8a4b8326 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2023-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 @@ -30,6 +30,7 @@ set(SRCS capacityScheduler.cpp createNewDecoderRequests.cpp contextProgress.cpp + contextTransferCoordinator.cpp dataTransceiver.cpp decoderBuffers.cpp kvCacheManager.cpp diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 2c69af0059be..e5a9c8c20273 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -36,6 +36,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" #include "tensorrt_llm/batch_manager/contextProgress.h" +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" #include "tensorrt_llm/batch_manager/dataTransceiver.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/kvCacheType.h" @@ -56,6 +57,7 @@ #include #include #include +#include #include #include @@ -521,6 +523,34 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mCacheSender = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); mCacheReceiver = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); + // Keep automatic enablement within the currently qualified C++ NIXL/UCX TP1/CP1 pipeline topology. + bool const coordinatorTopologyEligible = worldConfig.getPipelineParallelism() > 1 && useMPI() + && backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL + && common::getEnvNixlBackend() == "UCX" && worldConfig.getTensorParallelism() == 1 + && worldConfig.getContextParallelism() == 1 && !mCacheState->getParallelConfig().mEnableAttentionDP; + if (worldConfig.getPipelineParallelism() > 1 && useMPI()) + { + TLLM_CHECK(mGroupPipeParaComm != nullptr); + constexpr std::uint64_t kCoordinatorProtocolVersion = 1; + std::uint64_t const localVersion = coordinatorTopologyEligible ? kCoordinatorProtocolVersion : 0; + constexpr bool cancellationEnabled = false; + std::uint64_t const localProtocolMode = (localVersion << 1) | static_cast(cancellationEnabled); + std::vector protocolModes(static_cast(mGroupPipeParaComm->getSize())); + mGroupPipeParaComm->allgather(&localProtocolMode, protocolModes.data(), 1, mpi::MpiType::kUINT64); + TLLM_CHECK_WITH_INFO(std::all_of(protocolModes.begin(), protocolModes.end(), + [&](std::uint64_t const mode) { return mode == localProtocolMode; }), + "Context-transfer consensus protocol version or cancellation mode differs across PP ranks."); + if (localVersion != 0) + { + mContextTransferCoordinator = std::make_unique(mGroupPipeParaComm); + TLLM_LOG_INFO( + "Enable asynchronous context-transfer consensus version %llu for PP group of size %d; in-flight " + "cancellation=%s.", + static_cast(kCoordinatorProtocolVersion), mGroupPipeParaComm->getSize(), + cancellationEnabled ? "enabled" : "disabled"); + } + } + initializeCommState(); } @@ -530,6 +560,7 @@ CacheTransceiver::~CacheTransceiver() // plugin are still alive. The workers can access both during termination. mCacheSender.reset(); mCacheReceiver.reset(); + mContextTransferCoordinator.reset(); if (mWrapperLibHandle) { @@ -789,6 +820,17 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( toCompleteIdSet.insert(request->mRequestId); } + auto recordOutcome + = [&](RequestIdType const requestId, std::shared_ptr const& request, bool const failed) + { + recordLocalTransferOutcome(requestId, request, failed, mCompletedSenderRequestIds, mFailedSenderRequestIds, + mSenderRequestsAwaitingConsensus); + if (mContextTransferCoordinator) + { + mContextTransferCoordinator->publishLocalOutcome(requestId, failed); + } + }; + // Record local terminal outcomes for requests selected this round. The // request is reported only after all ranks in the sync group agree that the // request reached a terminal state. @@ -811,16 +853,16 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } if (blockAll || (toCompleteIdSet.find(requestId) != toCompleteIdSet.end())) { + bool terminal = false; + bool failed = false; try { auto const status = blockAll ? std::future_status::ready : future.wait_for(futureWaitInterval); if (status == std::future_status::ready) { future.get(); - bool const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; - recordLocalTransferOutcome(requestId, request, failed, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; + terminal = true; } else if (status == std::future_status::timeout) { @@ -834,18 +876,28 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { TLLM_LOG_ERROR( "Future returned unexpected status for request %ld. Recording as failed.", requestId); - - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + failed = true; + terminal = true; } } catch (std::exception const& e) { TLLM_LOG_ERROR("Error occurred during context transfer for request %ld: %s", requestId, e.what()); - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + failed = true; + terminal = true; + } + catch (...) + { + TLLM_LOG_ERROR("Unknown error occurred during context transfer for request %ld", requestId); + failed = true; + terminal = true; + } + if (terminal) + { + auto terminalRequest = request; it = mSenderFutures.erase(it); + // Publish outside the transfer-future try/catch. A protocol error must not rewrite an immutable vote. + recordOutcome(requestId, terminalRequest, failed); } } else @@ -855,8 +907,35 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } RequestStatuses requestsStatus{}; - auto const consensusOutcome - = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, mFailedSenderRequestIds); + TransferConsensusOutcome consensusOutcome; + if (mContextTransferCoordinator) + { + auto mergeCoordinatorOutcome = [&]() + { + auto coordinatorOutcome = mContextTransferCoordinator->poll(); + consensusOutcome.completedRequestIds.insert( + coordinatorOutcome.completedRequestIds.begin(), coordinatorOutcome.completedRequestIds.end()); + consensusOutcome.failedRequestIds.insert( + coordinatorOutcome.failedRequestIds.begin(), coordinatorOutcome.failedRequestIds.end()); + }; + do + { + mergeCoordinatorOutcome(); + if (blockAll + && consensusOutcome.completedRequestIds.size() + consensusOutcome.failedRequestIds.size() + < mSenderRequestsAwaitingConsensus.size()) + { + std::this_thread::yield(); + } + } while (blockAll + && consensusOutcome.completedRequestIds.size() + consensusOutcome.failedRequestIds.size() + < mSenderRequestsAwaitingConsensus.size()); + } + else + { + consensusOutcome + = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, mFailedSenderRequestIds); + } for (auto const requestId : consensusOutcome.failedRequestIds) { auto const requestIt = mSenderRequestsAwaitingConsensus.find(requestId); diff --git a/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp b/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp new file mode 100644 index 000000000000..909b7858c574 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp @@ -0,0 +1,475 @@ +/* + * 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. + */ + +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" + +#include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/logger.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager +{ + +ContextTransferVoteReducer::ContextTransferVoteReducer(int const participantCount) + : mParticipantCount(participantCount) +{ + TLLM_CHECK_WITH_INFO(participantCount > 0, "Context-transfer consensus requires at least one participant."); +} + +void ContextTransferVoteReducer::recordVote( + int const participantRank, std::uint64_t const requestId, ContextTransferVote const vote) +{ + TLLM_CHECK_WITH_INFO(participantRank >= 0 && participantRank < mParticipantCount, + "Context-transfer consensus participant rank is out of range."); + TLLM_CHECK_WITH_INFO(vote == ContextTransferVote::kCompleted || vote == ContextTransferVote::kFailed, + "Context-transfer consensus received an invalid vote."); + + auto [requestIt, inserted] = mRequestVotes.try_emplace(requestId, mParticipantCount); + static_cast(inserted); + auto& requestVotes = requestIt->second; + auto& recordedVote = requestVotes.votes.at(static_cast(participantRank)); + auto const packedVote = static_cast(vote); + if (recordedVote != 0) + { + TLLM_CHECK_WITH_INFO( + recordedVote == packedVote, "Context-transfer participant changed its terminal vote for a request."); + return; + } + + recordedVote = packedVote; + ++requestVotes.terminalCount; + requestVotes.failed = requestVotes.failed || vote == ContextTransferVote::kFailed; +} + +void ContextTransferVoteReducer::recordTimeout(std::uint64_t const requestId) +{ + auto [requestIt, inserted] = mRequestVotes.try_emplace(requestId, mParticipantCount); + static_cast(inserted); + auto& requestVotes = requestIt->second; + if (!requestVotes.timedOut) + { + requestVotes.timedOut = true; + requestVotes.timeoutPending = true; + } +} + +ContextTransferConsensusResult ContextTransferVoteReducer::takeReady() +{ + ContextTransferConsensusResult result; + for (auto requestIt = mRequestVotes.begin(); requestIt != mRequestVotes.end();) + { + auto& requestVotes = requestIt->second; + if (requestVotes.timeoutPending) + { + result.timedOutRequestIds.insert(requestIt->first); + requestVotes.timeoutPending = false; + } + if (requestVotes.terminalCount != mParticipantCount) + { + ++requestIt; + continue; + } + + auto& terminalRequestIds + = (requestVotes.failed || requestVotes.timedOut) ? result.failedRequestIds : result.completedRequestIds; + terminalRequestIds.insert(requestIt->first); + requestIt = mRequestVotes.erase(requestIt); + } + return result; +} + +void ContextTransferVoteReducer::clear() noexcept +{ + mRequestVotes.clear(); +} + +class ContextTransferCoordinator::Impl +{ +public: + explicit Impl(std::shared_ptr comm) + : mComm(std::move(comm)) + , mCoordinatorRank(0) + , mReducer(mComm ? mComm->getSize() : 1) + { + TLLM_CHECK_WITH_INFO(mComm != nullptr, "Context-transfer coordination requires a communicator."); + TLLM_CHECK_WITH_INFO(mComm->isMpi(), "Asynchronous context-transfer coordination requires MPI."); + TLLM_CHECK_WITH_INFO( + mComm->getSize() > 1, "Asynchronous context-transfer coordination requires multiple participants."); + mCoordinatorRank = mComm->getSize() - 1; + } + + void publishLocalOutcome(std::uint64_t const requestId, bool const failed) + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot publish a context-transfer vote after coordinator shutdown."); + auto const vote = failed ? ContextTransferVote::kFailed : ContextTransferVote::kCompleted; + auto const [voteIt, inserted] = mPublishedLocalVotes.emplace(requestId, vote); + TLLM_CHECK_WITH_INFO( + inserted || voteIt->second == vote, "This rank changed its terminal vote for a context transfer."); + if (!inserted) + { + return; + } + + try + { + if (isCoordinator()) + { + mReducer.recordVote(mComm->getRank(), requestId, vote); + } + else + { + queuePacket( + requestId, static_cast(vote), mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + } + catch (...) + { + mPublishedLocalVotes.erase(voteIt); + throw; + } + } + + void publishTimeout(std::uint64_t const requestId) + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot publish a context-transfer timeout after coordinator shutdown."); + TLLM_CHECK_WITH_INFO(mPublishedLocalVotes.find(requestId) == mPublishedLocalVotes.end(), + "A context-transfer timeout must be published before its immutable terminal vote."); + auto const [timeoutIt, inserted] = mPublishedTimeouts.insert(requestId); + if (!inserted) + { + return; + } + + try + { + if (isCoordinator()) + { + mReducer.recordTimeout(requestId); + } + else + { + queuePacket(requestId, kTimedOutSignal, mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + } + catch (...) + { + mPublishedTimeouts.erase(timeoutIt); + throw; + } + } + + [[nodiscard]] ContextTransferConsensusResult poll() + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot poll context-transfer coordination after shutdown."); + return progress(); + } + + void shutdown() noexcept + { + if (mShutdown) + { + return; + } + mShutdown = true; + + try + { + if (!isCoordinator()) + { + queuePacket(/*requestId=*/0, kCloseMarker, mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + + auto const deadline = std::chrono::steady_clock::now() + kShutdownTimeout; + while (!shutdownComplete()) + { + static_cast(progress()); + if (!shutdownComplete()) + { + if (std::chrono::steady_clock::now() >= deadline) + { + TLLM_LOG_ERROR( + "Timed out shutting down asynchronous context-transfer coordinator; rank=%d " + "peer_closes=%zu/%d ack=%d pending_sends=%zu. Aborting to avoid freeing active MPI " + "requests.", + mComm->getRank(), mClosedPeers.size(), mComm->getSize() - 1, mCloseAcknowledged, + mPendingSends.size()); + std::abort(); + } + std::this_thread::yield(); + } + } + } + catch (std::exception const& error) + { + TLLM_LOG_ERROR("Failed to shut down asynchronous context-transfer coordinator: %s", error.what()); + std::abort(); + } + catch (...) + { + TLLM_LOG_ERROR("Failed to shut down asynchronous context-transfer coordinator with an unknown error."); + std::abort(); + } + } + +private: + static constexpr std::size_t kPacketFieldCount = 2; + static constexpr std::uint64_t kTimedOutSignal = 3; + static constexpr std::uint64_t kCloseMarker = 4; + static constexpr auto kShutdownTimeout = std::chrono::seconds(30); + + struct PendingSend + { + std::array packet{}; + std::unique_ptr request; + }; + + [[nodiscard]] bool isCoordinator() const + { + return mComm->getRank() == mCoordinatorRank; + } + + void queuePacket(std::uint64_t const requestId, std::uint64_t const value, int const peer, mpi::MpiTag const tag) + { + mPendingSends.emplace_back(); + auto& pendingSend = mPendingSends.back(); + pendingSend.packet = {requestId, value}; + try + { + pendingSend.request = mComm->sendAsync( + pendingSend.packet.data(), pendingSend.packet.size(), mpi::MpiType::kUINT64, peer, tag); + } + catch (...) + { + mPendingSends.pop_back(); + throw; + } + } + + void queueUpdateForPeers(std::uint64_t const requestId, std::uint64_t const update) + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer != mCoordinatorRank) + { + queuePacket(requestId, update, peer, mpi::MpiTag::kContextTransferUpdate); + } + } + } + + void reapCompletedSends() + { + for (auto sendIt = mPendingSends.begin(); sendIt != mPendingSends.end();) + { + TLLM_CHECK(sendIt->request); + if (sendIt->request->isCompleted()) + { + sendIt = mPendingSends.erase(sendIt); + } + else + { + ++sendIt; + } + } + } + + void drainVotes() + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer == mCoordinatorRank) + { + continue; + } + + MPI_Status status{}; + while (mComm->iprobe(peer, mpi::MpiTag::kContextTransferEvent, &status)) + { + std::array packet{}; + mComm->recv( + packet.data(), packet.size(), mpi::MpiType::kUINT64, peer, mpi::MpiTag::kContextTransferEvent); + if (packet.back() == kCloseMarker) + { + TLLM_CHECK_WITH_INFO( + mClosedPeers.insert(peer).second, "Received a duplicate context-transfer close marker."); + continue; + } + TLLM_CHECK_WITH_INFO(mClosedPeers.find(peer) == mClosedPeers.end(), + "Received a context-transfer vote after its peer close marker."); + if (packet.back() == kTimedOutSignal) + { + mReducer.recordTimeout(packet.front()); + continue; + } + mReducer.recordVote(peer, packet.front(), static_cast(packet.back())); + } + } + } + + ContextTransferConsensusResult completeCoordinatorUpdates() + { + auto result = mReducer.takeReady(); + for (auto const requestId : result.timedOutRequestIds) + { + queueUpdateForPeers(requestId, kTimedOutSignal); + } + for (auto const requestId : result.failedRequestIds) + { + queueUpdateForPeers(requestId, static_cast(ContextTransferVote::kFailed)); + mPublishedLocalVotes.erase(requestId); + mPublishedTimeouts.erase(requestId); + } + for (auto const requestId : result.completedRequestIds) + { + queueUpdateForPeers(requestId, static_cast(ContextTransferVote::kCompleted)); + mPublishedLocalVotes.erase(requestId); + mPublishedTimeouts.erase(requestId); + } + return result; + } + + ContextTransferConsensusResult drainUpdates() + { + ContextTransferConsensusResult result; + MPI_Status status{}; + while (mComm->iprobe(mCoordinatorRank, mpi::MpiTag::kContextTransferUpdate, &status)) + { + std::array packet{}; + mComm->recv(packet.data(), packet.size(), mpi::MpiType::kUINT64, mCoordinatorRank, + mpi::MpiTag::kContextTransferUpdate); + if (packet.back() == kCloseMarker) + { + TLLM_CHECK_WITH_INFO(!mCloseAcknowledged, "Received a duplicate coordinator close marker."); + mCloseAcknowledged = true; + mPublishedLocalVotes.clear(); + mPublishedTimeouts.clear(); + continue; + } + + if (packet.back() == kTimedOutSignal) + { + result.timedOutRequestIds.insert(packet.front()); + continue; + } + + auto const outcome = static_cast(packet.back()); + TLLM_CHECK_WITH_INFO(outcome == ContextTransferVote::kCompleted || outcome == ContextTransferVote::kFailed, + "Received an invalid context-transfer commit outcome."); + auto const localVoteIt = mPublishedLocalVotes.find(packet.front()); + TLLM_CHECK_WITH_INFO(localVoteIt != mPublishedLocalVotes.end(), + "Received a context-transfer commit before publishing the local terminal vote."); + if (outcome == ContextTransferVote::kFailed) + { + result.failedRequestIds.insert(packet.front()); + } + else + { + result.completedRequestIds.insert(packet.front()); + } + mPublishedLocalVotes.erase(localVoteIt); + mPublishedTimeouts.erase(packet.front()); + } + return result; + } + + ContextTransferConsensusResult progress() + { + reapCompletedSends(); + if (isCoordinator()) + { + drainVotes(); + auto result = completeCoordinatorUpdates(); + if (mShutdown && !mCloseSent && mClosedPeers.size() == static_cast(mComm->getSize() - 1)) + { + // Shutdown is an explicit abort epoch. A peer close proves that no more votes will arrive from that + // peer, so incomplete requests cannot reach a global decision and are intentionally abandoned. + mReducer.clear(); + mPublishedLocalVotes.clear(); + mPublishedTimeouts.clear(); + for (int peer = 0; peer < mCoordinatorRank; ++peer) + { + queuePacket(/*requestId=*/0, kCloseMarker, peer, mpi::MpiTag::kContextTransferUpdate); + } + mCloseSent = true; + } + return result; + } + return drainUpdates(); + } + + [[nodiscard]] bool shutdownComplete() const + { + if (isCoordinator()) + { + return mCloseSent && mPendingSends.empty(); + } + return mCloseAcknowledged && mPendingSends.empty(); + } + + std::shared_ptr mComm; + int mCoordinatorRank; + ContextTransferVoteReducer mReducer; + std::unordered_map mPublishedLocalVotes; + std::unordered_set mPublishedTimeouts; + std::unordered_set mClosedPeers; + std::list mPendingSends; + bool mShutdown{false}; + bool mCloseSent{false}; + bool mCloseAcknowledged{false}; +}; + +ContextTransferCoordinator::ContextTransferCoordinator(std::shared_ptr comm) + : mImpl(std::make_unique(std::move(comm))) +{ +} + +ContextTransferCoordinator::~ContextTransferCoordinator() +{ + shutdown(); +} + +void ContextTransferCoordinator::publishLocalOutcome(std::uint64_t const requestId, bool const failed) +{ + mImpl->publishLocalOutcome(requestId, failed); +} + +void ContextTransferCoordinator::publishTimeout(std::uint64_t const requestId) +{ + mImpl->publishTimeout(requestId); +} + +ContextTransferConsensusResult ContextTransferCoordinator::poll() +{ + return mImpl->poll(); +} + +void ContextTransferCoordinator::shutdown() noexcept +{ + if (mImpl) + { + mImpl->shutdown(); + } +} + +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index eeb5ec34deef..575f031467fe 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -20,6 +20,7 @@ add_gtest(cacheTransBufferTest cacheTransBufferTest.cpp) add_gtest(bufferIndexHolderTest bufferIndexHolderTest.cpp) add_gtest(capacitySchedulerTest capacitySchedulerTest.cpp) add_gtest(contextProgressTest contextProgressTest.cu) +add_gtest(contextTransferCoordinatorTest contextTransferCoordinatorTest.cpp) add_gtest(evictionPolicyTest evictionPolicyTest.cpp) add_gtest(kvCacheManagerTest kvCacheManagerTest.cpp) add_gtest(kvCacheUtilsTest kvCacheUtilsTest.cpp) diff --git a/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp b/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp new file mode 100644 index 000000000000..7b0f7aba3fc4 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp @@ -0,0 +1,143 @@ +/* + * 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. + */ + +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" + +#include + +namespace tensorrt_llm::batch_manager +{ +namespace +{ + +TEST(ContextTransferVoteReducerTest, WaitsForEveryParticipant) +{ + ContextTransferVoteReducer reducer(4); + reducer.recordVote(0, 17, ContextTransferVote::kCompleted); + reducer.recordVote(1, 17, ContextTransferVote::kCompleted); + reducer.recordVote(2, 17, ContextTransferVote::kCompleted); + EXPECT_TRUE(reducer.takeReady().completedRequestIds.empty()); + + reducer.recordVote(3, 17, ContextTransferVote::kCompleted); + auto const result = reducer.takeReady(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{17}); + EXPECT_TRUE(result.failedRequestIds.empty()); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, FailureWinsAfterEveryParticipantVotes) +{ + ContextTransferVoteReducer reducer(4); + reducer.recordVote(0, 23, ContextTransferVote::kCompleted); + reducer.recordVote(1, 23, ContextTransferVote::kFailed); + reducer.recordVote(2, 23, ContextTransferVote::kCompleted); + reducer.recordVote(3, 23, ContextTransferVote::kCompleted); + + auto const result = reducer.takeReady(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{23}); +} + +TEST(ContextTransferVoteReducerTest, RejectsChangedTerminalVote) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordVote(0, 29, ContextTransferVote::kCompleted); + reducer.recordVote(0, 29, ContextTransferVote::kCompleted); + EXPECT_ANY_THROW(reducer.recordVote(0, 29, ContextTransferVote::kFailed)); +} + +TEST(ContextTransferVoteReducerTest, AccumulatesInterleavedRequestsIndependently) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordVote(0, 31, ContextTransferVote::kCompleted); + reducer.recordVote(1, 37, ContextTransferVote::kFailed); + reducer.recordVote(1, 31, ContextTransferVote::kCompleted); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{31}); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(0, 37, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{37}); +} + +TEST(ContextTransferVoteReducerTest, RejectsInvalidInput) +{ + EXPECT_ANY_THROW(ContextTransferVoteReducer(0)); + ContextTransferVoteReducer reducer(2); + EXPECT_ANY_THROW(reducer.recordVote(-1, 41, ContextTransferVote::kCompleted)); + EXPECT_ANY_THROW(reducer.recordVote(2, 41, ContextTransferVote::kCompleted)); + EXPECT_ANY_THROW(reducer.recordVote(0, 41, static_cast(99))); +} + +TEST(ContextTransferVoteReducerTest, TimeoutIsVisibleBeforeTerminalFailure) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordTimeout(43); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{43}); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_TRUE(result.failedRequestIds.empty()); + + result = reducer.takeReady(); + EXPECT_TRUE(result.timedOutRequestIds.empty()); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(0, 43, ContextTransferVote::kCompleted); + reducer.recordVote(1, 43, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{43}); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, DuplicateTimeoutIsIdempotent) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordTimeout(47); + reducer.recordTimeout(47); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{47}); + + reducer.recordTimeout(47); + result = reducer.takeReady(); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, TimeoutAfterPartialTerminalVotesStillWins) +{ + ContextTransferVoteReducer reducer(3); + reducer.recordVote(2, 53, ContextTransferVote::kCompleted); + reducer.recordTimeout(53); + reducer.recordVote(0, 53, ContextTransferVote::kCompleted); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{53}); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(1, 53, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{53}); + EXPECT_TRUE(result.completedRequestIds.empty()); +} + +} // namespace +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index d1ca104cca1a..569d8738f24f 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -31,6 +31,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" @@ -42,6 +43,7 @@ #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" +#include #include #include #include @@ -53,6 +55,7 @@ #include #include #include +#include #include "gtest/gtest.h" #include @@ -89,6 +92,95 @@ T serializeDeserialize(T const& val) } // namespace +TEST(ContextTransferCoordinatorTest, CommitsStaggeredSuccessAndFailureWithoutCollectivePolling) +{ + auto& world = tensorrt_llm::mpi::MpiComm::world(); + if (world.getSize() < 2) + { + GTEST_SKIP() << "mpirun with at least two processes is required to run this test."; + } + + auto comm = std::make_shared(std::addressof(world)); + { + ContextTransferCoordinator coordinator(comm); + auto waitForOutcome = [&](std::uint64_t const requestId, bool const expectFailure) + { + bool observed = false; + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!observed && std::chrono::steady_clock::now() < deadline) + { + auto result = coordinator.poll(); + observed = expectFailure ? result.failedRequestIds.count(requestId) != 0 + : result.completedRequestIds.count(requestId) != 0; + if (!observed) + { + std::this_thread::yield(); + } + } + EXPECT_TRUE(observed); + }; + auto waitForTimeout = [&](std::uint64_t const requestId) + { + bool observed = false; + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!observed && std::chrono::steady_clock::now() < deadline) + { + auto const result = coordinator.poll(); + observed = result.timedOutRequestIds.count(requestId) != 0; + if (!observed) + { + std::this_thread::yield(); + } + } + EXPECT_TRUE(observed); + }; + + constexpr std::uint64_t kCompletedRequestId = 644815201; + world.barrier(); + if (world.getRank() == world.getSize() - 1) + { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + coordinator.publishLocalOutcome(kCompletedRequestId, /*failed=*/false); + if (world.getRank() != world.getSize() - 1) + { + auto const earlyResult = coordinator.poll(); + EXPECT_TRUE(earlyResult.completedRequestIds.empty()); + EXPECT_TRUE(earlyResult.failedRequestIds.empty()); + } + waitForOutcome(kCompletedRequestId, /*expectFailure=*/false); + + constexpr std::uint64_t kFailedRequestId = 644815202; + world.barrier(); + coordinator.publishLocalOutcome(kFailedRequestId, /*failed=*/world.getRank() == 0); + waitForOutcome(kFailedRequestId, /*expectFailure=*/true); + + constexpr std::uint64_t kTimedOutRequestId = 644815203; + world.barrier(); + if (world.getRank() == 0) + { + coordinator.publishTimeout(kTimedOutRequestId); + coordinator.publishTimeout(kTimedOutRequestId); + } + waitForTimeout(kTimedOutRequestId); + coordinator.publishLocalOutcome(kTimedOutRequestId, /*failed=*/false); + waitForOutcome(kTimedOutRequestId, /*expectFailure=*/true); + + constexpr std::uint64_t kInvalidOrderingRequestId = 644815204; + world.barrier(); + coordinator.publishLocalOutcome(kInvalidOrderingRequestId, /*failed=*/false); + EXPECT_ANY_THROW(coordinator.publishTimeout(kInvalidOrderingRequestId)); + waitForOutcome(kInvalidOrderingRequestId, /*expectFailure=*/false); + + // Exercise asymmetric but orderly teardown after every decision has committed. + if (world.getRank() == 0) + { + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } + } + world.barrier(); +} + class RequestInfoTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) { public: diff --git a/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp b/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp index 221cd98b5f02..b8596d511fab 100644 --- a/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -25,6 +25,8 @@ #endif // ENABLE_MULTI_DEVICE #include +#include +#include namespace mpi = tensorrt_llm::mpi; namespace tr = tensorrt_llm::runtime; @@ -47,6 +49,29 @@ TEST(MPIUtils, WorldRankAndSize) EXPECT_LE(rank, size); } +#if ENABLE_MULTI_DEVICE +TEST(MPIUtils, AsyncSendCanBePolled) +{ + auto& comm = mpi::MpiComm::world(); + auto const rank = comm.getRank(); + auto const size = comm.getSize(); + auto const destination = (rank + 1) % size; + auto const source = (rank + size - 1) % size; + std::uint64_t const sentValue = static_cast(rank); + std::uint64_t receivedValue = 0; + + auto request + = comm.sendAsync(&sentValue, 1, mpi::MpiType::kUINT64, destination, mpi::MpiTag::kContextTransferEvent); + static_cast(comm.recv(&receivedValue, 1, mpi::MpiType::kUINT64, source, mpi::MpiTag::kContextTransferEvent)); + while (!request->isCompleted()) + { + std::this_thread::yield(); + } + + EXPECT_EQ(receivedValue, static_cast(source)); +} +#endif // ENABLE_MULTI_DEVICE + template void testBroadcast() { From 3301daec44ef7db6b73b45f6969f4675c2407a79 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 24 Jul 2026 09:41:32 -0700 Subject: [PATCH 6/9] [NVBUG-6448152][test] restore native CMake source order Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> (cherry picked from commit 4b182f1118073faf43f7206808696c6fba02596e) --- cpp/tensorrt_llm/batch_manager/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index 154b8a4b8326..1389d0de068f 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -33,11 +33,11 @@ set(SRCS contextTransferCoordinator.cpp dataTransceiver.cpp decoderBuffers.cpp - kvCacheManager.cpp encoderBuffers.cpp guidedDecoder.cpp handleContextLogits.cpp handleGenerationLogits.cpp + kvCacheManager.cpp kvCacheEventManager.cpp kvCacheTransferManager.cpp kvCacheManagerV2Utils.cpp From 767e97f7373746dadd4775d9e8e6be4331c5d922 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:29:54 -0700 Subject: [PATCH 7/9] [NVBUG-6448152][test] apply current-toolchain build compatibility Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- cpp/include/tensorrt_llm/runtime/virtualMemory.h | 5 +---- cpp/tensorrt_llm/deep_ep/CMakeLists.txt | 15 +++++++++++++++ cpp/tensorrt_llm/runtime/virtualMemory.cpp | 11 ++--------- 3 files changed, 18 insertions(+), 13 deletions(-) diff --git a/cpp/include/tensorrt_llm/runtime/virtualMemory.h b/cpp/include/tensorrt_llm/runtime/virtualMemory.h index a7e95b42d707..d74673f37e01 100644 --- a/cpp/include/tensorrt_llm/runtime/virtualMemory.h +++ b/cpp/include/tensorrt_llm/runtime/virtualMemory.h @@ -505,10 +505,7 @@ class CudaVirtualMemoryAllocator { std::size_t gpuAlignment = 1; CUmemAllocationProp const prop{CU_MEM_ALLOCATION_TYPE_PINNED, CU_MEM_HANDLE_TYPE_NONE, - { - CU_MEM_LOCATION_TYPE_DEVICE, - device, - }}; + CUmemLocation{CU_MEM_LOCATION_TYPE_DEVICE, {device}}}; TLLM_CU_CHECK( cuMemGetAllocationGranularity(&gpuAlignment, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); alignment = std::lcm(getpagesize(), gpuAlignment); diff --git a/cpp/tensorrt_llm/deep_ep/CMakeLists.txt b/cpp/tensorrt_llm/deep_ep/CMakeLists.txt index 562c9e7d694c..e00815eed13e 100644 --- a/cpp/tensorrt_llm/deep_ep/CMakeLists.txt +++ b/cpp/tensorrt_llm/deep_ep/CMakeLists.txt @@ -120,6 +120,15 @@ if(NOT CMAKE_CXX_COMPILER_ID STREQUAL "GNU") set(CMAKE_C_COMPILER gcc) set(CMAKE_CXX_COMPILER g++) set(CMAKE_CUDA_HOST_COMPILER g++) + # PyTorch's cmake/public/cuda.cmake (loaded transitively by + # find_package(Torch)) appends -Xcompiler=-fclang-abi-compat=17 to + # CMAKE_CUDA_FLAGS whenever the parent build is configured with Clang>=18 (see + # pytorch PR #175233). Since this subdirectory falls back to GCC for NVSHMEM + # compatibility, that Clang-only flag would be forwarded to g++ via `nvcc + # -ccbin=g++` and abort the build with: g++: error: unrecognized command-line + # option '-fclang-abi-compat=17' + string(REGEX REPLACE "-Xcompiler=-fclang-abi-compat=[0-9]+" "" + CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}") endif() # Add nvshmem external project @@ -204,6 +213,12 @@ target_compile_options( target_compile_definitions( deep_ep_cpp_tllm PRIVATE DISABLE_AGGRESSIVE_PTX_INSTRS TORCH_EXTENSION_NAME=deep_ep_cpp_tllm) +# Newer CUDA containers provide NVSHMEM headers in the default CUDA include +# directory. DeepEP must compile against the vendored NVSHMEM headers because it +# links the vendored NVSHMEM static library below. +target_include_directories( + deep_ep_cpp_tllm BEFORE + PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/nvshmem-build/src/include) target_link_libraries( deep_ep_cpp_tllm PRIVATE nvshmem_project::nvshmem ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIB}) diff --git a/cpp/tensorrt_llm/runtime/virtualMemory.cpp b/cpp/tensorrt_llm/runtime/virtualMemory.cpp index 9b23866d6281..0d08012a29d8 100644 --- a/cpp/tensorrt_llm/runtime/virtualMemory.cpp +++ b/cpp/tensorrt_llm/runtime/virtualMemory.cpp @@ -344,11 +344,7 @@ void CudaVirtualMemoryAllocator::allocate(Pointer* ptr, std::size_t n, int devic CUDAVirtualMemoryChunk::Configurators configurators; configurators.push_back(std::make_unique(address, alignedSize, - CUmemAccessDesc{{ - CU_MEM_LOCATION_TYPE_DEVICE, - device, - }, - CU_MEM_ACCESS_FLAGS_PROT_READWRITE})); + CUmemAccessDesc{CUmemLocation{CU_MEM_LOCATION_TYPE_DEVICE, {device}}, CU_MEM_ACCESS_FLAGS_PROT_READWRITE})); switch (mConfig->mMode) { @@ -368,10 +364,7 @@ void CudaVirtualMemoryAllocator::allocate(Pointer* ptr, std::size_t n, int devic mConfig->mManager.add(address, mConfig->mTag, std::make_unique>(CUmemAllocationProp{CU_MEM_ALLOCATION_TYPE_PINNED, CU_MEM_HANDLE_TYPE_NONE, - { - CU_MEM_LOCATION_TYPE_DEVICE, - device, - }}, + CUmemLocation{CU_MEM_LOCATION_TYPE_DEVICE, {device}}}, alignedSize), std::move(configurators)); From 96266ce34f0d45870031e2fe92e108c7a9c3dbe0 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:27:11 -0700 Subject: [PATCH 8/9] [NVBUG 6448152][infra] align pre-admission test dependencies Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- requirements.txt | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/requirements.txt b/requirements.txt index 99fe8fcb8cbe..a243f9b26ab5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -21,14 +21,14 @@ pandas h5py==3.12.1 StrEnum sentencepiece>=0.1.99 -tensorrt~=10.15.1 -# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-02.html#rel-26-02 uses 2.11.0a0. -torch>=2.10.0,<=2.11.0a0 +tensorrt~=10.16.1 +# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-04.html#rel-26-04 uses 2.12.0a0. +torch>=2.11.0,<=2.12.0a0 torchvision nvidia-modelopt[torch]~=0.37.0 -# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-02.html#rel-26-02 uses 2.29.2 -# torch 2.10.0+cu130 depends on nvidia-nccl-cu13==2.28.9 -nvidia-nccl-cu13>=2.28.9,<=2.29.2 +# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-04.html#rel-26-04 uses 2.29.7 +# torch 2.11.0+cu130 depends on nvidia-nccl-cu13==2.28.9 +nvidia-nccl-cu13>=2.28.9,<=2.29.7 nvidia-cuda-nvrtc transformers==5.5.4 prometheus_client @@ -89,3 +89,5 @@ etcd-sdk-python==0.0.7 python-multipart smg-grpc-proto>=0.4.2 cache-dit>=1.3.5 +librosa +msgpack From f08296b2e61abff5c00bbe084b8bde54ee617909 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:21:17 -0700 Subject: [PATCH 9/9] [NVBUG 6448152][infra] filter clang flag for FlashMLA GCC Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- cpp/tensorrt_llm/flash_mla/CMakeLists.txt | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/cpp/tensorrt_llm/flash_mla/CMakeLists.txt b/cpp/tensorrt_llm/flash_mla/CMakeLists.txt index e87f12275f60..024cee521d10 100644 --- a/cpp/tensorrt_llm/flash_mla/CMakeLists.txt +++ b/cpp/tensorrt_llm/flash_mla/CMakeLists.txt @@ -44,6 +44,15 @@ if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64" AND CMAKE_CXX_COMPILER_ID set(CMAKE_CUDA_HOST_COMPILER ${GCC_EXECUTABLE}) message( STATUS "FlashMLA: Using GCC at ${GCC_EXECUTABLE} for CUDA compilation") + # PyTorch's cmake/public/cuda.cmake (loaded transitively by + # find_package(Torch)) appends -Xcompiler=-fclang-abi-compat=17 to + # CMAKE_CUDA_FLAGS whenever the parent build is configured with Clang>=18 (see + # pytorch PR #175233). Since CUDA host compilation here falls back to GCC, + # that Clang-only flag would be forwarded to g++ via `nvcc -ccbin=g++` and + # abort the build with: g++: error: unrecognized command-line option + # '-fclang-abi-compat=17' + string(REGEX REPLACE "-Xcompiler=-fclang-abi-compat=[0-9]+" "" + CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}") endif() # Check CUDA version and architecture support