diff --git a/inference/a3mega/llama3.1-70b/trtllm-gke/README.md b/inference/a3mega/llama3.1-70b/trtllm-gke/README.md new file mode 100644 index 00000000..01ca07d6 --- /dev/null +++ b/inference/a3mega/llama3.1-70b/trtllm-gke/README.md @@ -0,0 +1,392 @@ +# Single Host Model Serving with NVIDIA TensorRT-LLM (TRT-LLM) on A3mega GKE Node Pool + +This document outlines the steps to serve and benchmark various Large Language Models (LLMs) using the [NVIDIA TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) framework on a single [A3-Mega GKE Node pool](https://cloud.google.com/kubernetes-engine). + +This guide walks you through setting up the necessary cloud infrastructure, configuring your environment, and deploying a high-performance LLM for inference. + + +## Table of Contents + +* [1. Test Environment](#test-environment) +* [2. High-Level Architecture](#architecture) +* [3. Environment Setup (One-Time)](#environment-setup) + * [3.1. Clone the Repository](#clone-repo) + * [3.2. Configure Environment Variables](#configure-vars) + * [3.3. Connect to your GKE Cluster](#connect-cluster) + * [3.4. Get Hugging Face Token](#get-hf-token) + * [3.5. Create Hugging Face Kubernetes Secret](#setup-hf-secret) +* [4. Run the Recipe](#run-the-recipe) + * [4.1. Supported Models](#supported-models) + * [4.2. Deploy and Benchmark a Model](#deploy-model) +* [5. Monitoring and Troubleshooting](#monitoring) + * [5.1. Check Deployment Status](#check-status) + * [5.2. View Logs](#view-logs) +* [6. Cleanup](#cleanup) + + +## 1. Test Environment + +[Back to Top](#table-of-contents) + +The recipe uses the following setup: + +* **Orchestration**: [Google Kubernetes Engine (GKE)](https://cloud.google.com/kubernetes-engine) +* **Deployment Configuration**: A [Helm chart](https://helm.sh/) is used to configure and deploy a [Kubernetes Deployment](https://kubernetes.io/docs/concepts/workloads/controllers/deployment/). This deployment encapsulates the inference of the target LLM using the TensorRT-LLM framework. + +This recipe has been optimized for and tested with the following configuration: + +* **GKE Cluster**: + * A [regional standard cluster](https://cloud.google.com/kubernetes-engine/docs/concepts/configuration-overview) version: `1.33.4-gke.1036000` or later. + * A GPU node pool with 1 [a3-megagpu-8g](https://docs.cloud.google.com/compute/docs/accelerator-optimized-machines#a3-mega-vms) machine. + * [Workload Identity Federation for GKE](https://cloud.google.com/kubernetes-engine/docs/concepts/workload-identity) enabled. + * [Cloud Storage FUSE CSI driver for GKE](https://cloud.google.com/kubernetes-engine/docs/concepts/cloud-storage-fuse-csi-driver) enabled. + * [DCGM metrics](https://cloud.google.com/kubernetes-engine/docs/how-to/dcgm-metrics) enabled. + * [Kueue](https://kueue.sigs.k8s.io/docs/reference/kueue.v1beta1/) and [JobSet](https://jobset.sigs.k8s.io/docs/overview/) APIs installed. + * Kueue configured to support [Topology Aware Scheduling](https://kueue.sigs.k8s.io/docs/concepts/topology_aware_scheduling/). +* A regional Google Cloud Storage (GCS) bucket to store logs generated by the recipe runs. + +> [!IMPORTANT] +> To prepare the required environment, see the [GKE environment setup guide](../../../../docs/configuring-environment-gke-a3-mega.md). +> Provisioning a new GKE cluster is a long-running operation and can take **20-30 minutes**. + + +## 2. High-Level Flow + +[Back to Top](#table-of-contents) + +Here is a simplified diagram of the flow that we follow in this recipe: + +```mermaid +--- +config: + layout: dagre +--- +flowchart TD + subgraph workstation["Client Workstation"] + T["Cluster Toolkit"] + B("Kubernetes API") + A["helm install"] + end + subgraph huggingface["Hugging Face Hub"] + I["Model Weights"] + end + subgraph gke["GKE Cluster (A3-Mega)"] + C["Deployment"] + D["Pod"] + E["TensorRT-LLM container"] + F["Service"] + end + subgraph storage["Cloud Storage"] + J["Bucket"] + end + + %% Logical/actual flow + T -- Create Cluster --> gke + A --> B + B --> C & F + C --> D + D --> E + F --> C + E -- Downloads at runtime --> I + E -- Write logs --> J + + + %% Layout control + gke +``` + +* **helm:** A package manager for Kubernetes to define, install, and upgrade applications. It's used here to configure and deploy the Kubernetes Deployment. +* **Deployment:** Manages the lifecycle of your model server pod, ensuring it stays running. +* **Service:** Provides a stable network endpoint (a DNS name and IP address) to access your model server. +* **Pod:** The smallest deployable unit in Kubernetes. The Triton server container with TensorRT-LLM runs inside this pod on a GPU-enabled node. +* **Cloud Storage:** A Cloud Storage bucket to store benchmark logs and other artifacts. + + +## 3. Environment Setup (One-Time) + +[Back to Top](#table-of-contents) + +First, you'll configure your local environment. These steps are required once before you can deploy any models. + + +### 3.1. Clone the Repository + +```bash +git clone https://github.com/ai-hypercomputer/gpu-recipes.git +cd gpu-recipes +export REPO_ROOT=$(pwd) +export RECIPE_ROOT=$REPO_ROOT/inference/a3mega/llama3.1-70b/trtllm-gke +``` + + +### 3.2. Configure Environment Variables + +This is the most critical step. These variables are used in subsequent commands to target the correct resources. + +```bash +export PROJECT_ID= +export CLUSTER_REGION= +export CLUSTER_NAME= +export KUEUE_NAME= +export GCS_BUCKET= +export TRTLLM_VERSION=1.3.0rc3 + +# Set the project for gcloud commands +gcloud config set project $PROJECT_ID +``` + +Replace the following values: + +| Variable | Description | Example | +| --------------------- | ------------------------------------------------------------------------------------------------------- | ------------------------------------------------------- | +| `PROJECT_ID` | Your Google Cloud Project ID. | `gcp-project-12345` | +| `CLUSTER_REGION` | The GCP region where your GKE cluster is located. | `us-central1` | +| `CLUSTER_NAME` | The name of your GKE cluster. | `a3-mega` | +| `KUEUE_NAME` | The name of the Kueue local queue. The default queue created by the cluster toolkit is `a3mega`. Verify the name in your cluster. | `a3mega` | +| `ARTIFACT_REGISTRY` | Full path to your Artifact Registry repository. | `us-central1-docker.pkg.dev/gcp-project-12345/my-repo` | +| `GCS_BUCKET` | Name of your GCS bucket (do not include `gs://`). | `my-benchmark-logs-bucket` | +| `TRTLLM_VERSION` | The tag/version for the Docker image. Other verions can be found at [NGC Catalog](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release) | `1.3.0rc3` | + + + +### 3.3. Connect to your GKE Cluster + +Fetch credentials for `kubectl` to communicate with your cluster. + +```bash +gcloud container clusters get-credentials $CLUSTER_NAME --region $CLUSTER_REGION +``` + + +### 3.4. Get Hugging Face token + +To access models through Hugging Face, you'll need a Hugging Face token. + 1. Create a [Hugging Face account](https://huggingface.co/) if you don't have one. + 2. For **gated models** like Llama 4, ensure you have requested and been granted access on Hugging Face before proceeding. + 3. Generate an Access Token: Go to **Your Profile > Settings > Access Tokens**. + 4. Select **New Token**. + 5. Specify a Name and a Role of at least `Read`. + 6. Select **Generate a token**. + 7. Copy the generated token to your clipboard. You'll use this later. + + + +### 3.5. Create Hugging Face Kubernetes Secret + +Create a Kubernetes Secret with your Hugging Face token to enable the pod to download model checkpoints from Hugging Face. + +```bash +# Paste your Hugging Face token here +export HF_TOKEN= + +kubectl create secret generic hf-secret \ +--from-literal=hf_api_token=${HF_TOKEN} \ +--dry-run=client -o yaml | kubectl apply -f - +``` + + +## 4. Run the recipe + +[Back to Top](#table-of-contents) + +> [!NOTE] +> After running the recipe with `helm install`, it can take **up to 30 minutes** for the deployment to become fully available. This is because the GKE node must first pull the Docker image and then download the model weights from Hugging Face. + +> [!TIP] +> You can use the [NVIDIA Model Optimizer](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_ptq) to quantize these models to FP8. + + +### 4.1. Supported Models + +[Back to Top](#table-of-contents) + +This recipe supports the following models. Running TRTLLM inference benchmarking on these models are only tested and validated on A3-Mega GKE nodes with certain combination of TP, PP, EP, number of GPU chips, input & output sequence length, precision, etc. + +As the PyTorch backend requires pre-quantized models for optimal performance, we use the FP8 quantized version for Llama 3.1 70B. + +| Model Name | Hugging Face ID | Configuration File | Release Name Suffix | +| :--- | :--- | :--- | :--- | +| **Llama 3.1 70B (FP8)** | `nvidia/Llama-3.1-70B-Instruct-FP8` | `llama3.1-70b.yaml` | `llama-3-1-70b` | + +> [!TIP] +> You can use the NVIDIA Model Optimizer to quantize these models to FP8. + + +### 4.2. Deploy and Benchmark a Model + +[Back to Top](#table-of-contents) + +The recipe uses [`trtllm-bench`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/legacy/performance/perf-benchmarking.md), a command-line tool from NVIDIA to benchmark the performance of TensorRT-LLM engine. + +1. **Configure model-specific variables.** Choose a model from the [table above](#supported-models) and set the variables: + + ```bash + # Example for Llama 3.1 70B (FP8) + export HF_MODEL_ID="nvidia/Llama-3.1-70B-Instruct-FP8" + export CONFIG_FILE="llama3.1-70b.yaml" + export RELEASE_NAME="$USER-llama-3-1-70b" + ``` + +2. **Install the helm chart.** You can run a single benchmark configuration or a sequence of multiple experiments by indexing the `experiments` list. + + ```bash + cd $RECIPE_ROOT + helm install -f values.yaml \ + --set workload.benchmarks.experiments[0].isl=128 \ + --set workload.benchmarks.experiments[0].osl=128 \ + --set workload.benchmarks.experiments[0].num_requests=1000 \ + --set-file workload_launcher=$REPO_ROOT/src/launchers/trtllm-launcher.sh \ + --set-file serving_config=$REPO_ROOT/src/frameworks/a3mega/trtllm-configs/${CONFIG_FILE} \ + --set queue=${KUEUE_NAME} \ + --set "volumes.gcsMounts[0].bucketName=${GCS_BUCKET}" \ + --set workload.model.name=${HF_MODEL_ID} \ + --set workload.image=nvcr.io/nvidia/tensorrt-llm/release:${TRTLLM_VERSION} \ + --set workload.framework=trtllm \ + ${RELEASE_NAME} \ + $REPO_ROOT/src/helm-charts/a3mega/trtllm-inference/single-node + ``` + > [!NOTE] + > You can modify the benchmark configuration at runtime by changing the values for `isl`, `osl`, and `num_requests` (number of prompts) in the Helm command to test different scenarios. + +3. **Check the deployment status:** + + ```bash + kubectl get deployment/${RELEASE_NAME}-serving + ``` + + Wait until the `READY` column shows `1/1`. See the [Monitoring and Troubleshooting](#monitoring) section to view the deployment logs. + + +## 5. Monitoring and Troubleshooting + +[Back to Top](#table-of-contents) + +After the model is deployed via Helm as described in the sections [above](#run-the-recipe), use the following steps to monitor the deployment and interact with the model. Replace `` and `` with the appropriate names from the model-specific deployment instructions (e.g., `$USER-llama3.1-70b` and `$USER-llama3.1-70b-svc`). + + +### 5.1. Check Deployment Status + +Check the status of your deployment. Replace the name if you deployed a different model. + +```bash +kubectl get deployment/$RELEASE_NAME-serving +``` + +Wait until the `READY` column shows `1/1`. If it shows `0/1`, the pod is still starting up. + +> [!NOTE] +> In the GKE UI on Cloud Console, you might see a status of "Does not have minimum availability" during startup. This is normal and will resolve once the pod is ready. + + +### 5.2. View Logs + +To see the logs from the TRTLLM server (useful for debugging), use the `-f` flag to follow the log stream: + +```bash +kubectl logs -f deployment/$RELEASE_NAME-serving +``` + +You should see logs indicating preparing the model, and then running the throughput benchmark test, similar to this: + +```bash +Running benchmark for nvidia/Llama3.1-70b with ISL=128, OSL=128, TP=8, EP=1, PP=1 + +=========================================================== +PYTORCH BACKEND +=========================================================== +Model: nvidia/Llama-3.1-70B-Instruct-FP8 +Model Path: /ssd/nvidia/Llama-3.1-70B-Instruct-FP8 +TensorRT LLM Version: 1.2 +Dtype: bfloat16 +KV Cache Dtype: FP8 +Quantization: FP8 + +=========================================================== +REQUEST DETAILS +=========================================================== +Number of requests: 1000 +Number of concurrent requests: 126.1051 +Average Input Length (tokens): 128.0000 +Average Output Length (tokens): 128.0000 +=========================================================== +WORLD + RUNTIME INFORMATION +=========================================================== +TP Size: 8 +PP Size: 1 +EP Size: 1 +Max Runtime Batch Size: 2304 +Max Runtime Tokens: 4608 +Scheduling Policy: GUARANTEED_NO_EVICT +KV Memory Percentage: 85.00% +Issue Rate (req/sec): 8.3913E+13 + +=========================================================== +PERFORMANCE OVERVIEW +=========================================================== +Request Throughput (req/sec): X.XX +Total Output Throughput (tokens/sec): X.XX +Total Token Throughput (tokens/sec): X.XX +Total Latency (ms): X.XX +Average request latency (ms): X.XX +Per User Output Throughput [w/ ctx] (tps/user): X.XX +Per GPU Output Throughput (tps/gpu): X.XX + +-- Request Latency Breakdown (ms) ----------------------- + +[Latency] P50 : X.XX +[Latency] P90 : X.XX +[Latency] P95 : X.XX +[Latency] P99 : X.XX +[Latency] MINIMUM: X.XX +[Latency] MAXIMUM: X.XX +[Latency] AVERAGE: X.XX + +=========================================================== +DATASET DETAILS +=========================================================== +Dataset Path: /ssd/token-norm-dist_llama3.1-70b_128_128_tp8.json +Number of Sequences: 1000 + +-- Percentiles statistics --------------------------------- + + Input Output Seq. Length +----------------------------------------------------------- +MIN: 128.0000 128.0000 256.0000 +MAX: 128.0000 128.0000 256.0000 +AVG: 128.0000 128.0000 256.0000 +P50: 128.0000 128.0000 256.0000 +P90: 128.0000 128.0000 256.0000 +P95: 128.0000 128.0000 256.0000 +P99: 128.0000 128.0000 256.0000 +=========================================================== +``` + + +## 6. Cleanup + +To avoid incurring further charges, clean up the resources you created. + +1. **Uninstall the Helm Release:** + + First, list your releases to get the deployed models: + + ```bash + # list deployed models + helm list --filter $USER + ``` + + Then, uninstall the desired release: + + ```bash + helm uninstall + ``` + Replace `` with the helm release names listed. + +2. **Delete the Kubernetes Secret:** + + ```bash + kubectl delete secret hf-secret --ignore-not-found=true + ``` + +3. (Optional) Clean up files in your GCS bucket if benchmarking was performed. +4. (Optional) Delete the [test environment](#test-environment) provisioned including GKE cluster. diff --git a/inference/a3mega/llama3.1-70b/trtllm-gke/values.yaml b/inference/a3mega/llama3.1-70b/trtllm-gke/values.yaml new file mode 100644 index 00000000..b2f9c233 --- /dev/null +++ b/inference/a3mega/llama3.1-70b/trtllm-gke/values.yaml @@ -0,0 +1,66 @@ +# Copyright 2026 Google LLC +# +# 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. + +targetPlatform: "gke" +clusterName: "" +queue: "" + +huggingface: + secretName: "hf-secret" + secretData: + token: "hf_api_token" + +workload: + framework: "trtllm" + gpus: 8 + image: "" + model: + name: "" + configPath: "/workload/configs" + configFile: "serving-args.yaml" + + benchmarks: + experiments: + - isl: 128 + osl: 128 + num_requests: 1000 + concurrency: 128 + max_seq_len: 32768 + +volumes: + gcsVolumes: true + ssdMountPath: "/ssd" + gcsMounts: + - bucketName: "" + mountPath: "/gcs" + +gpuPlatformSettings: + useHostPlugin: false + ncclPluginImage: "us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/nccl-plugin-gpudirecttcpx-dev:v1.0.15" + rxdmImage: "us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpgpudmarxd-dev:v1.0.21" + ncclBuildType: "228" + +network: + hostNetwork: true + ncclSettings: + - name: NCCL_DEBUG + value: "WARN" + subnetworks: [] + +trtllm: + replicaCount: 1 + service: + type: ClusterIP + ports: + http: 8000 diff --git a/src/frameworks/a3mega/trtllm-configs/llama3.1-70b.yaml b/src/frameworks/a3mega/trtllm-configs/llama3.1-70b.yaml new file mode 100644 index 00000000..bd1d6a12 --- /dev/null +++ b/src/frameworks/a3mega/trtllm-configs/llama3.1-70b.yaml @@ -0,0 +1,4 @@ +tp_size: 8 +pp_size: 1 +kv_cache_free_gpu_mem_fraction: 0.85 +backend: pytorch \ No newline at end of file diff --git a/src/helm-charts/a3mega/trtllm-inference/single-node/Chart.yaml b/src/helm-charts/a3mega/trtllm-inference/single-node/Chart.yaml new file mode 100644 index 00000000..4bebfbfc --- /dev/null +++ b/src/helm-charts/a3mega/trtllm-inference/single-node/Chart.yaml @@ -0,0 +1,20 @@ +# Copyright 2026 Google LLC +# +# 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. + +apiVersion: v2 +name: trtllm-llama-3-1-70b-inference +description: A Helm chart for running TensorRT-LLM inference on a single A3 Mega GKE node for llama-3-1-70b. +type: application +version: 0.1.0 +appVersion: "1.0.0" \ No newline at end of file diff --git a/src/helm-charts/a3mega/trtllm-inference/single-node/templates/model-serve-launcher.yaml b/src/helm-charts/a3mega/trtllm-inference/single-node/templates/model-serve-launcher.yaml new file mode 100644 index 00000000..187dce64 --- /dev/null +++ b/src/helm-charts/a3mega/trtllm-inference/single-node/templates/model-serve-launcher.yaml @@ -0,0 +1,417 @@ +# Copyright 2026 Google LLC +# +# 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. + +{{- /* Extract benchmark parameters from Helm values */ -}} +{{- $experiment := index .Values.workload.benchmarks.experiments 0 -}} +{{- $isl := $experiment.isl -}} +{{- $osl := $experiment.osl -}} +{{- $concurrency := $experiment.concurrency | default "" -}} +{{- $max_seq_len := $experiment.max_seq_len | default "32768" -}} +{{- $num_requests := $experiment.num_requests -}} + +# --- CONFIGMAPS --- +apiVersion: v1 +kind: ConfigMap +metadata: + name: "{{ .Release.Name }}-config" +data: + serving-configuration: |- +{{- if .Values.serving_config }} +{{ .Values.serving_config | nindent 4 }} +{{- else }} + tp_size: {{ .Values.workload.model.tp_size | default 8 }} + pp_size: {{ .Values.workload.model.pp_size | default 1 }} + quantization: {{ .Values.workload.quantization | default "FP8" }} + kv_cache_free_gpu_mem_fraction: {{ .Values.workload.kv_cache_free_gpu_mem_fraction | default 0.70 }} + backend: {{ .Values.workload.backend | default "tensorrt" }} +{{- end }} +--- +apiVersion: v1 +kind: ConfigMap +metadata: + name: "{{ .Release.Name }}-launcher" +data: + launch-workload.sh: |- +{{- if .Values.workload_launcher }} +{{ .Values.workload_launcher | nindent 4 }} +{{- else }} + #!/bin/bash + echo "No workload launcher specified" + exit 1 +{{- end }} + +--- + +{{ $nodes := div .Values.workload.gpus 8 | max 1 }} +{{ $gpusPerNode := min .Values.workload.gpus 8 }} + +{{- $root := . -}} + +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ .Release.Name }}-serving + labels: + app: {{ .Release.Name }}-serving +spec: + replicas: {{ $nodes }} + selector: + matchLabels: + app: {{ .Release.Name }}-serving + template: + metadata: + labels: + app: {{ .Release.Name }}-serving + annotations: + checksum/serving-config: {{ .Values.serving_config | default "" | sha256sum }} + checksum/launcher-config: {{ .Values.workload_launcher | default "" | sha256sum }} + kubectl.kubernetes.io/default-container: serving + gke-gcsfuse/volumes: "true" + gke-gcsfuse/cpu-limit: "0" + gke-gcsfuse/memory-limit: "0" + gke-gcsfuse/ephemeral-storage-limit: "0" + devices.gke.io/container.tcpxo-daemon: |+ + - path: /dev/nvidia0 + - path: /dev/nvidia1 + - path: /dev/nvidia2 + - path: /dev/nvidia3 + - path: /dev/nvidia4 + - path: /dev/nvidia5 + - path: /dev/nvidia6 + - path: /dev/nvidia7 + - path: /dev/nvidiactl + - path: /dev/nvidia-uvm + - path: /dev/dmabuf_import_helper + networking.gke.io/default-interface: "eth0" + spec: + restartPolicy: Always + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + subdomain: "{{.Release.Name}}" + {{ if $root.Values.targetNodes }} + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + - key: kubernetes.io/hostname + operator: In + values: + {{- range $hostname := $root.Values.targetNodes }} + - {{ $hostname }} + {{- end }} + {{ end }} + + tolerations: + - key: nvidia.com/gpu + operator: Exists + - key: cloud.google.com/impending-node-termination + operator: Exists + volumes: + - name: nvidia-dir-host + hostPath: + path: /home/kubernetes/bin/nvidia + + - name: serving-configuration + configMap: + name: "{{ .Release.Name }}-config" + items: + - key: "serving-configuration" + path: "serving-args.yaml" + - name: serving-launcher + configMap: + name: "{{ .Release.Name }}-launcher" + defaultMode: 0700 + + {{- if not $root.Values.gpuPlatformSettings.useHostPlugin }} + - name: nccl-plugin-volume + emptyDir: {} + {{- end }} + - name: sys + hostPath: + path: /sys + - name: proc-sys + hostPath: + path: /proc/sys + - name: aperture-devices + hostPath: + path: /dev/aperture_devices + - name: workload-terminated-volume + emptyDir: {} + - name: local-ssd + hostPath: + path: /mnt/stateful_partition/kube-ephemeral-ssd + - name: shared-memory + emptyDir: + medium: "Memory" + sizeLimit: 250Gi + {{- if and $root.Values.volumes (hasKey $root.Values.volumes "gcsMounts") }} + {{- range $gcs := $root.Values.volumes.gcsMounts }} + - name: "{{ $gcs.bucketName }}" + csi: + driver: gcsfuse.csi.storage.gke.io + volumeAttributes: + bucketName: "{{ $gcs.bucketName }}" + {{- end}} + {{- end}} + + initContainers: + {{- if not $root.Values.gpuPlatformSettings.useHostPlugin }} + - name: nccl-plugin-installer + image: "{{ $root.Values.gpuPlatformSettings.ncclPluginImage }}" + imagePullPolicy: Always + volumeMounts: + - name: nccl-plugin-volume + mountPath: /usr/local/tcpxo + env: + - name: BUILD_TYPE + value: "{{ $root.Values.gpuPlatformSettings.ncclBuildType }}" + command: + - bash + - -c + - | + set -ex + chmod 755 /scripts/container_entry.sh + /scripts/container_entry.sh install --install-nccl --nccl-buildtype ${BUILD_TYPE} + cp -r /var/lib/tcpxo/* /usr/local/tcpxo/ + {{- end }} + + - name: tcpxo-daemon + image: {{ $root.Values.gpuPlatformSettings.rxdmImage }} + imagePullPolicy: Always + securityContext: + capabilities: + add: + - NET_ADMIN + - NET_BIND_SERVICE + restartPolicy: Always + volumeMounts: + - name: nvidia-dir-host + mountPath: /usr/local/nvidia + - name: sys + mountPath: /hostsysfs + - name: proc-sys + mountPath: /hostprocsysfs + env: + - name: LD_LIBRARY_PATH + value: /usr/local/nvidia/lib64 + command: + - bash + - -c + - | + cleanup() { + echo "Received SIGTERM, exiting RxDM" + if [ -n "$child_pid" ]; then + echo "Sending SIGTERM to child process" + kill -TERM "$child_pid" + fi + exit 0 + } + trap cleanup SIGTERM + + chmod 755 /fts/entrypoint_rxdm_container.sh + /fts/entrypoint_rxdm_container.sh --num_hops=2 --num_nics=8 --uid= --alsologtostderr & child_pid=$! + + wait "$child_pid" + + containers: + - name: serving + {{- if typeIs "string" $root.Values.workload.image }} + image: "{{ $root.Values.workload.image }}" + {{- else }} + image: "{{ $root.Values.workload.image.repository }}:{{ $root.Values.workload.image.tag }}" + {{- end }} + securityContext: + privileged: true + resources: + requests: + nvidia.com/gpu: {{ $gpusPerNode }} + limits: + nvidia.com/gpu: {{ $gpusPerNode }} + env: + - name: JOB_ORCHESTRATOR + value: "gke" + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: "{{ $root.Values.huggingface.secretName }}" + key: "{{ $root.Values.huggingface.secretData.token }}" + - name: HF_HUB_ENABLE_HF_TRANSFER + value: "0" + - name: HF_HOME + value: "/ssd" + - name: MODEL_DOWNLOAD_DIR + value: "/ssd/{{ $root.Values.workload.model.name }}" + - name: MODEL_NAME + value: "{{ $root.Values.workload.model.name }}" + - name: LAUNCHER_SCRIPT + value: "/workload/launcher/launch-workload.sh" + - name: SERVER_ARGS_FILE + value: "/workload/configs/serving-args.yaml" + - name: TRTLLM_DIR + value: "/app/tensorrt_llm" + - name: TORCH_DISTRIBUTED_DEBUG + value: "INFO" + - name: GLOO_SOCKET_IFNAME + value: "eth0" + + {{- if $root.Values.gpuPlatformSettings.useHostPlugin }} + - name: LD_LIBRARY_PATH + value: "/usr/local/nvidia/lib64" + - name: NCCL_LIB_DIR + value: "/usr/local/nvidia/lib64" + {{- else }} + - name: LD_LIBRARY_PATH + value: "/usr/local/tcpxo/lib64:/usr/local/cuda/lib64:/usr/local/nvidia/lib64" + {{- end }} + - name: LD_PRELOAD + value: "/usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2" + - name: NCCL_NET_PLUGIN_PATH + value: "/usr/local/tcpxo/lib64" + - name: NCCL_P2P_LEVEL + value: "NVL" + - name: NCCL_FASTRAK_LLCM_DEVICE_DIRECTORY + value: /dev/aperture_devices + - name: NCCL_FASTRAK_CTRL_DEV + value: "eth0" + - name: NCCL_SOCKET_IFNAME + value: "eth0" + - name: NCCL_CROSS_NIC + value: "0" + - name: NCCL_ALGO + value: "Ring,Tree" + - name: NCCL_PROTO + value: "Simple,LL128" + - name: NCCL_MIN_NCHANNELS + value: "4" + - name: NCCL_DYNAMIC_CHUNK_SIZE + value: "524288" + - name: NCCL_P2P_NET_CHUNKSIZE + value: "524288" + - name: NCCL_P2P_PCI_CHUNKSIZE + value: "524288" + - name: NCCL_P2P_NVL_CHUNKSIZE + value: "1048576" + - name: NCCL_FASTRAK_NUM_FLOWS + value: "2" + - name: NCCL_BUFFSIZE + value: "8388608" + - name: NCCL_NET_GDR_LEVEL + value: "PIX" + - name: NCCL_DEBUG_SUBSYS + value: "INIT,GRAPH,ENV,TUNING,NET" + - name: NCCL_FASTRAK_ENABLE_HOTPATH_LOGGING + value: "0" + - name: NCCL_FASTRAK_USE_SNAP + value: "1" + - name: NCCL_FASTRAK_ENABLE_CONTROL_CHANNEL + value: "0" + - name: NCCL_FASTRAK_USE_LLCM + value: "1" + - name: NCCL_TUNER_PLUGIN + value: "libnccl-tuner.so" + - name: NCCL_TUNER_CONFIG_PATH + value: "/usr/local/tcpxo/lib64/a3plus_tuner_config.textproto" + - name: NCCL_SHIMNET_GUEST_CONFIG_CHECKER_CONFIG_FILE + value: "/usr/local/tcpxo/lib64/a3plus_guest_config.textproto" + - name: NCCL_FASTRAK_PLUGIN_ACCEPT_TIMEOUT_MS + value: "600000" + - name: CUDA_VISIBLE_DEVICES + value: "0,1,2,3,4,5,6,7" + - name: NCCL_FASTRAK_IFNAME + value: "eth1,eth2,eth3,eth4,eth5,eth6,eth7,eth8" + - name: NCCL_NVLSTREE_MAX_CHUNKSIZE + value: "131072" + - name: NVTE_FWD_LAYERNORM_SM_MARGIN + value: "8" + - name: NVTE_BWD_LAYERNORM_SM_MARGIN + value: "8" + - name: NCCL_P2P_PXN_LEVEL + value: "0" + - name: NCCL_NET_PLUGIN_TELEMETRY_MODE + value: "1" + - name: NCCL_GPUVIZ_ENABLE_MILLISECOND_BANDWIDTH_OUTPUT + value: "1" + - name: NCCL_GPUVIZ_FILE_ROTATION_INTERVAL_IN_SECONDS + value: "300" + - name: TLLM_NUMA_AWARE_WORKER_AFFINITY + value: "1" + + {{- range $environment_variable := $root.Values.network.ncclSettings }} + - name: {{ $environment_variable.name }} + value: "{{ $environment_variable.value }}" + {{- end }} + + command: + - bash + - -c + - | + + # Performance Tuning and NCCL Fixes for A3 Mega + ulimit -l unlimited + export NCCL_NET_PLUGIN=none + export NCCL_TUNER_PLUGIN=none + + trtllm-bench() { + unset NCCL_P2P_LEVEL + command trtllm-bench "$@" --concurrency {{ $concurrency }} --max_seq_len {{ $max_seq_len }} + } + export -f trtllm-bench + + ldconfig 2>/dev/null + chmod +x "$LAUNCHER_SCRIPT" + + ARGS=() + while IFS=': ' read -r key value || [ -n "$key" ]; do + [[ -z "$key" || "$key" == \#* ]] && continue + key=$(echo "$key" | xargs); value=$(echo "$value" | xargs) + if [[ "$value" == "true" ]]; then ARGS+=("--$key"); elif [[ "$value" == "false" ]]; then ARGS+=("--$key" "false"); elif [ -n "$value" ]; then ARGS+=("--$key" "$value"); else ARGS+=("--$key"); fi + done < "/workload/configs/serving-args.yaml" + + echo "Starting benchmark..." + exec "$LAUNCHER_SCRIPT" --model_name "{{ .Values.workload.model.name }}" --isl "{{ $isl }}" --osl "{{ $osl }}" --num_requests "{{ $num_requests }}" -- "${ARGS[@]}" + echo "-----------------------------------------------------------" + echo "Benchmarks complete." + ports: + - containerPort: {{ $root.Values.trtllm.service.ports.http }} + + volumeMounts: + - name: nvidia-dir-host + mountPath: /usr/local/nvidia + - name: aperture-devices + mountPath: /dev/aperture_devices + {{- if not $root.Values.gpuPlatformSettings.useHostPlugin }} + - name: nccl-plugin-volume + mountPath: /usr/local/tcpxo + {{- end }} + - name: sys + mountPath: /hostsysfs + - name: proc-sys + mountPath: /hostprocsysfs + - name: workload-terminated-volume + mountPath: /semaphore + - name: shared-memory + mountPath: /dev/shm + - name: local-ssd + mountPath: "{{ $root.Values.volumes.ssdMountPath }}" + - name: serving-configuration + mountPath: {{ $root.Values.workload.configPath | default "/workload/configs" }} + - name: serving-launcher + mountPath: /workload/launcher + {{- if and $root.Values.volumes (hasKey $root.Values.volumes "gcsMounts") }} + {{- range $gcs := $root.Values.volumes.gcsMounts }} + - name: "{{ $gcs.bucketName }}" + mountPath: "{{ $gcs.mountPath }}" + {{- end }} + {{- end }} diff --git a/src/helm-charts/a3mega/trtllm-inference/single-node/templates/model-serve-svc.yaml b/src/helm-charts/a3mega/trtllm-inference/single-node/templates/model-serve-svc.yaml new file mode 100644 index 00000000..5566e370 --- /dev/null +++ b/src/helm-charts/a3mega/trtllm-inference/single-node/templates/model-serve-svc.yaml @@ -0,0 +1,26 @@ +# Copyright 2026 Google LLC +# +# 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. + +apiVersion: v1 +kind: Service +metadata: + name: {{ .Release.Name }}-svc +spec: + selector: + app: {{ .Release.Name }}-serving + ports: + - name: http + port: {{ .Values.trtllm.service.ports.http }} + targetPort: {{ .Values.trtllm.service.ports.http }} + type: {{ .Values.trtllm.service.type }} \ No newline at end of file diff --git a/src/launchers/trtllm-launcher.sh b/src/launchers/trtllm-launcher.sh index dc3f828b..dbdd6cfc 100644 --- a/src/launchers/trtllm-launcher.sh +++ b/src/launchers/trtllm-launcher.sh @@ -213,7 +213,7 @@ run_benchmark() { --ep $ep_size \ --backend "pytorch" \ --kv_cache_free_gpu_mem_fraction $kv_cache_free_gpu_mem_fraction \ - $extra_args $vl_args > $output_file + $extra_args $vl_args | tee "$output_file" else echo "Building engine" trtllm-bench \ @@ -234,10 +234,9 @@ run_benchmark() { --model_path /ssd/${model_name} throughput \ --dataset $dataset_file \ --engine_dir $engine_dir \ - --kv_cache_free_gpu_mem_fraction $kv_cache_free_gpu_mem_fraction $extra_args >$output_file + --kv_cache_free_gpu_mem_fraction $kv_cache_free_gpu_mem_fraction $extra_args | tee $output_file fi - cat $output_file gcloud storage cp $output_file /gcs/benchmark_logs/trtllm/ rm -rf $engine_dir