feat(generation): add TensorRT-LLM rollout backend

.agents/contributor-skills/build-and-dependency/SKILL.md

@@ -26,10 +26,11 @@ docker buildx build -f docker/Dockerfile \
 Skip optional backends to reduce build time:
 
 ```bash
-# Skip vLLM and SGLang
+# Skip vLLM, SGLang, and TRT-LLM
 docker buildx build -f docker/Dockerfile \
     --build-arg SKIP_VLLM_BUILD=1 \
     --build-arg SKIP_SGLANG_BUILD=1 \
+    --build-arg SKIP_TRTLLM_BUILD=1 \
     --tag nemo-rl:latest .
 ```
 

.github/workflows/_build_container.yml

@@ -35,6 +35,16 @@ on:
         default: ""
         description: Additional Docker build contexts.
         type: string
+      trtllm-ccache-tag:
+        required: false
+        default: ""
+        description: Prefix for the TRT-LLM ccache image tag; the target architecture and a hash of the effective base image are appended automatically. Empty disables the ccache.
+        type: string
+      trtllm-wheel-cache-tag:
+        required: false
+        default: ""
+        description: Prefix for the TRT-LLM wheel cache image tag; the target architecture and a hash of the effective base image are appended automatically. Empty disables the cache.
+        type: string
       dockerfile:
         required: true
         description: Path to the Dockerfile.
@@ -83,6 +93,142 @@ jobs:
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
 
+      - name: Resolve TRT-LLM cache configuration
+        id: trtllm_cache
+        shell: bash
+        env:
+          BUILD_ARGS: ${{ inputs.build-args }}
+          CCACHE_TAG_PREFIX: ${{ inputs.trtllm-ccache-tag }}
+          WHEEL_TAG_PREFIX: ${{ inputs.trtllm-wheel-cache-tag }}
+          DOCKERFILE: ${{ inputs.dockerfile }}
+          PLATFORM: ${{ inputs.platform }}
+        run: |
+          set -euo pipefail
+
+          # Match Dockerfile semantics: any non-empty value skips TRT-LLM.
+          SKIP_VALUE=$(printf '%s\n' "$BUILD_ARGS" | sed -nE 's/^[[:space:]]*SKIP_TRTLLM_BUILD[[:space:]]*=(.*)$/\1/p' | tail -1)
+          SKIP_VALUE=$(printf '%s' "$SKIP_VALUE" | sed -E 's/^[[:space:]]+//; s/[[:space:]]+$//')
+          if [[ -n "$SKIP_VALUE" ]]; then
+            echo "TRT-LLM build is disabled by SKIP_TRTLLM_BUILD"
+            echo "enabled=false" >> "$GITHUB_OUTPUT"
+            exit 0
+          fi
+          echo "enabled=true" >> "$GITHUB_OUTPUT"
+
+          # No cache configuration is needed when both image prefixes are empty.
+          if [[ -z "$CCACHE_TAG_PREFIX" && -z "$WHEEL_TAG_PREFIX" ]]; then
+            exit 0
+          fi
+
+          BASE_IMAGE=$(printf '%s\n' "$BUILD_ARGS" | sed -nE 's/^[[:space:]]*BASE_IMAGE[[:space:]]*=[[:space:]]*(.+)[[:space:]]*$/\1/p' | tail -1)
+          if [[ -z "$BASE_IMAGE" ]]; then
+            BASE_IMAGE=$(sed -nE 's/^[[:space:]]*ARG[[:space:]]+BASE_IMAGE[[:space:]]*=[[:space:]]*(.+)[[:space:]]*$/\1/p' "$DOCKERFILE" | tail -1)
+          fi
+          if [[ -z "$BASE_IMAGE" ]]; then
+            echo "Could not resolve BASE_IMAGE from build args or $DOCKERFILE" >&2
+            exit 1
+          fi
+
+          ARCH="${PLATFORM#*/}"
+          ARCH="${ARCH%%/*}"
+          if [[ -z "$ARCH" || "$PLATFORM" != */* || "$PLATFORM" == *,* ]]; then
+            echo "Expected a single OS/architecture platform, got: $PLATFORM" >&2
+            exit 1
+          fi
+
+          CACHE_KEY=$(printf '%s' "$BASE_IMAGE" | sha256sum | cut -c1-12)
+          {
+            echo "arch=$ARCH"
+            if [[ -n "$CCACHE_TAG_PREFIX" ]]; then
+              echo "ccache_tag=${CCACHE_TAG_PREFIX}-${ARCH}-${CACHE_KEY}"
+            fi
+            if [[ -n "$WHEEL_TAG_PREFIX" ]]; then
+              echo "wheel_tag=${WHEEL_TAG_PREFIX}-${ARCH}-${CACHE_KEY}"
+            fi
+          } >> "$GITHUB_OUTPUT"
+          echo "TRT-LLM caches use $BASE_IMAGE on $ARCH (key $CACHE_KEY)"
+
+      - name: Resolve TRT-LLM cache seeds
+        id: trtllm_cache_seed
+        if: steps.trtllm_cache.outputs.ccache_tag != '' || steps.trtllm_cache.outputs.wheel_tag != ''
+        shell: bash
+        env:
+          CCACHE_TAG: ${{ steps.trtllm_cache.outputs.ccache_tag }}
+          WHEEL_TAG: ${{ steps.trtllm_cache.outputs.wheel_tag }}
+        run: |
+          set -euo pipefail
+
+          resolve_seed() {
+            local name="$1"
+            local tag="$2"
+            local context="$3"
+            local description="$4"
+            [[ -n "$tag" ]] || return 0
+
+            local ref="$REGISTRY/$IMAGE_NAME:$tag"
+            if docker buildx imagetools inspect "$ref" >/dev/null 2>&1; then
+              echo "Using TRT-LLM $description seed $ref"
+              {
+                echo "${name}_ref=$ref"
+                echo "${name}_context=${context}=docker-image://$ref"
+              } >> "$GITHUB_OUTPUT"
+            else
+              echo "No TRT-LLM $description image found at $ref; starting with an empty cache"
+            fi
+          }
+
+          resolve_seed ccache "$CCACHE_TAG" trtllm-ccache-seed ccache
+          resolve_seed wheel "$WHEEL_TAG" trtllm-wheel-cache-seed "wheel cache"
+
+      - name: Configure TRT-LLM cache exports
+        id: trtllm_cache_export
+        if: steps.trtllm_cache.outputs.ccache_tag != '' || steps.trtllm_cache.outputs.wheel_tag != ''
+        shell: bash
+        env:
+          CCACHE_SEED_REF: ${{ steps.trtllm_cache_seed.outputs.ccache_ref }}
+          CCACHE_TAG: ${{ steps.trtllm_cache.outputs.ccache_tag }}
+          WHEEL_SEED_REF: ${{ steps.trtllm_cache_seed.outputs.wheel_ref }}
+          WHEEL_TAG: ${{ steps.trtllm_cache.outputs.wheel_tag }}
+          PLATFORM: ${{ inputs.platform }}
+        run: |
+          set -euo pipefail
+
+          configure_export() {
+            local name="$1"
+            local seed_ref="$2"
+            local output_context="$3"
+            local description="$4"
+            local max_incremental_layers=16
+            local cache_layout=incremental-v1
+            local export_mode=full
+
+            if [[ -n "$seed_ref" ]]; then
+              local image_metadata
+              image_metadata=$(docker buildx imagetools inspect "$seed_ref" --format "{{len (index .Image \"$PLATFORM\").RootFS.DiffIDs}} {{index (index .Image \"$PLATFORM\").Config.Labels \"com.nvidia.nemo-rl.cache-layout\"}}" 2>/dev/null || true)
+              if [[ -z "$image_metadata" ]]; then
+                image_metadata=$(docker buildx imagetools inspect "$seed_ref" --format '{{len .Image.RootFS.DiffIDs}} {{index .Image.Config.Labels "com.nvidia.nemo-rl.cache-layout"}}' 2>/dev/null || true)
+              fi
+
+              local layer_count existing_layout
+              read -r layer_count existing_layout <<< "$image_metadata"
+              if [[ "$existing_layout" == "$cache_layout" && "$layer_count" =~ ^[0-9]+$ ]] && (( layer_count < max_incremental_layers )); then
+                export_mode=incremental
+                echo "Appending an incremental $description layer to $seed_ref ($layer_count existing layers)"
+                echo "${name}_output_context=${output_context}=docker-image://$seed_ref" >> "$GITHUB_OUTPUT"
+              else
+                echo "Compacting TRT-LLM $description into a full snapshot (layout: ${existing_layout:-legacy}, layers: ${layer_count:-unknown})"
+              fi
+            fi
+            echo "${name}_mode=$export_mode" >> "$GITHUB_OUTPUT"
+          }
+
+          if [[ -n "$CCACHE_TAG" ]]; then
+            configure_export ccache "$CCACHE_SEED_REF" trtllm-ccache-output-base ccache
+          fi
+          if [[ -n "$WHEEL_TAG" ]]; then
+            configure_export wheel "$WHEEL_SEED_REF" trtllm-wheel-cache-output-base "wheel cache"
+          fi
+
       - name: Get recently merged PR cache refs
         id: recent_pr_cache_refs
         uses: actions/github-script@v8
@@ -160,7 +306,10 @@ jobs:
           push: true
           context: .
           platforms: ${{ inputs.platform }}
-          build-contexts: ${{ inputs.build-contexts }}
+          build-contexts: |
+            ${{ inputs.build-contexts }}
+            ${{ steps.trtllm_cache_seed.outputs.ccache_context }}
+            ${{ steps.trtllm_cache_seed.outputs.wheel_context }}
           build-args: ${{ inputs.build-args }}
           cache-from: |
             ${{ steps.build_meta.outputs.cache-from }}
@@ -170,3 +319,56 @@ jobs:
           tags: |
             ${{ steps.build_meta.outputs.tags }}
           target: ${{ inputs.target }}
+
+      - name: Update TRT-LLM ccache image
+        # Cache mounts survive a failed RUN in the current Buildx builder. Run
+        # this step after both successful and failed image builds so completed
+        # compiler outputs are preserved for the next retry.
+        if: ${{ always() && !cancelled() && steps.trtllm_cache.outputs.enabled == 'true' && steps.trtllm_cache.outputs.ccache_tag != '' && steps.trtllm_cache_export.outcome == 'success' }}
+        uses: docker/build-push-action@v5
+        with:
+          file: ${{ inputs.dockerfile }}
+          push: true
+          context: .
+          platforms: ${{ inputs.platform }}
+          build-contexts: |
+            ${{ inputs.build-contexts }}
+            ${{ steps.trtllm_cache_seed.outputs.ccache_context }}
+            ${{ steps.trtllm_cache_export.outputs.ccache_output_context }}
+          build-args: |
+            ${{ inputs.build-args }}
+            CCACHE_EXPORT_NONCE=${{ github.run_id }}-${{ github.run_attempt }}
+            CCACHE_EXPORT_MODE=${{ steps.trtllm_cache_export.outputs.ccache_mode }}
+          cache-from: |
+            ${{ steps.build_meta.outputs.cache-from }}
+            ${{ steps.recent_pr_cache_refs.outputs.cache-from }}
+          no-cache: false
+          provenance: false
+          tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ steps.trtllm_cache.outputs.ccache_tag }}
+          target: trtllm-ccache
+
+      - name: Update TRT-LLM wheel cache image
+        # Persist the custom wheel even when the release build fails, so a
+        # retry can reuse completed compilation work.
+        if: ${{ always() && !cancelled() && steps.trtllm_cache.outputs.enabled == 'true' && steps.trtllm_cache.outputs.wheel_tag != '' && steps.trtllm_cache_export.outcome == 'success' }}
+        uses: docker/build-push-action@v5
+        with:
+          file: ${{ inputs.dockerfile }}
+          push: true
+          context: .
+          platforms: ${{ inputs.platform }}
+          build-contexts: |
+            ${{ inputs.build-contexts }}
+            ${{ steps.trtllm_cache_seed.outputs.wheel_context }}
+            ${{ steps.trtllm_cache_export.outputs.wheel_output_context }}
+          build-args: |
+            ${{ inputs.build-args }}
+            WHEEL_CACHE_EXPORT_NONCE=${{ github.run_id }}-${{ github.run_attempt }}
+            WHEEL_CACHE_EXPORT_MODE=${{ steps.trtllm_cache_export.outputs.wheel_mode }}
+          cache-from: |
+            ${{ steps.build_meta.outputs.cache-from }}
+            ${{ steps.recent_pr_cache_refs.outputs.cache-from }}
+          no-cache: false
+          provenance: false
+          tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ steps.trtllm_cache.outputs.wheel_tag }}
+          target: trtllm-wheel-cache

.github/workflows/cicd-main.yml

@@ -359,8 +359,11 @@ jobs:
       build-contexts: |
         nemo-rl=.
         ${{ vars.UV_BUILD_CACHE == 'enabled' && format('uv-cache-seed=docker-image://{0}/{1}:uv-cache', needs.org-member-pre-flight.outputs.registry, vars.CI_CONTAINER_NAME) || '' }}
+      trtllm-ccache-tag: ${{ vars.TRTLLM_BUILD_CACHE == 'enabled' && 'trtllm-ccache' || '' }}
+      trtllm-wheel-cache-tag: ${{ vars.TRTLLM_BUILD_CACHE == 'enabled' && 'trtllm-wheel-cache' || '' }}
       build-args: |
         MAX_JOBS=4
+        TRTLLM_BUILD_JOBS=24
         NEMO_RL_COMMIT=${{ needs.pre-flight.outputs.test_sha }}
 
   build-container-gb200:
@@ -398,8 +401,11 @@ jobs:
       build-contexts: |
         nemo-rl=.
         ${{ vars.UV_BUILD_CACHE == 'enabled' && format('uv-cache-seed=docker-image://{0}/{1}:uv-cache', needs.gb200-config.outputs.registry, vars.CI_CONTAINER_NAME) || '' }}
+      trtllm-ccache-tag: ${{ vars.TRTLLM_BUILD_CACHE == 'enabled' && 'trtllm-ccache' || '' }}
+      trtllm-wheel-cache-tag: ${{ vars.TRTLLM_BUILD_CACHE == 'enabled' && 'trtllm-wheel-cache' || '' }}
       build-args: |
         MAX_JOBS=4
+        TRTLLM_BUILD_JOBS=8
         NEMO_RL_COMMIT=${{ needs.pre-flight.outputs.test_sha }}
 
   update-uv-cache:
@@ -522,6 +528,8 @@ jobs:
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L0_Unit_Tests_Sglang
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
+          - script: L0_Unit_Tests_Trtllm
+            runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L0_Unit_Tests_Mcore
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L0_Unit_Tests_Mcore_Policy_1
@@ -582,6 +590,9 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v6
       - name: main
+        # Lfast reuses the main image. Until TRT-LLM lands on main, running this
+        # shard would make prefetch_venvs compile the cp313 wheel from source.
+        if: ${{ needs.pre-flight.outputs.test_level != 'Lfast' || matrix.script != 'L0_Unit_Tests_Trtllm' }}
         uses: ./.github/actions/test-template
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
@@ -706,6 +717,8 @@ jobs:
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L1_Functional_Tests_SGLang
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
+          - script: L1_Functional_Tests_Trtllm
+            runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L1_Functional_Tests_Gym
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L1_Functional_Tests_GRPO_1