Adds a sample inferencing recipe for TRT-LLM on A4X (#50)

* WIP recipe for a4x trtllm inference

* Plumbs through additional properties needed to support setting llm_api_args, customizing kvcache free mem, and adds an actual output example.

* Delete lightweight test model config

Author: Chris113113

inference/a4x/single-host-serving/tensorrt-llm/README.md

@@ -0,0 +1,63 @@
+# Copyright 2025 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.
+
+queue:
+
+dwsSettings:
+  maxRunDurationSeconds:
+
+huggingface:
+  secretName: hf-secret
+  secretData:
+    token: "hf_api_token"
+
+volumes:
+  gcsVolumes: true
+  ssdMountPath: "/ssd"
+  gcsMounts:
+    - bucketName:
+      mountPath: "/gcs"
+
+service:
+  type: ClusterIP
+  ports:
+    http: 8000
+
+workload:
+  model:
+    name:
+  gpus: 4
+  image:
+  framework:
+  configFile: serving-args.yaml
+  configPath: /workload/configs
+  envs:
+    - name: HF_HUB_ENABLE_HF_TRANSFER
+      value: "1"
+    - name: LAUNCHER_SCRIPT
+      value: "/workload/launcher/launch-workload.sh"
+    - name: SERVER_ARGS_FILE
+      value: "/workload/configs/serving-args.yaml"
+  benchmarks:
+    experiments:
+      - isl: 128
+        osl: 128
+        num_requests: 1000
+
+network:
+  subnetworks[]:
+  gibVersion: us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib-arm64:v1.0.7
+  ncclSettings:
+    - name: NCCL_DEBUG
+      value: "VERSION"

inference/a4x/single-host-serving/tensorrt-llm/values.yaml

@@ -0,0 +1,35 @@
+tp_size: 4
+ep_size: 4
+pp_size: 1
+backend: pytorch
+kv_cache_free_gpu_mem_fraction: 0.85
+llm_api_args:
+  cuda_graph_config:
+    batch_sizes:
+    - 1
+    - 2
+    - 4
+    - 8
+    - 16
+    - 20
+    - 24
+    - 32
+    - 64
+    - 96
+    - 128
+    - 160
+    - 192
+    - 256
+    - 320
+    - 384
+    - 512
+    enable_padding: true
+  enable_attention_dp: true
+  enable_chunked_prefill: true
+  kv_cache_config:
+    dtype: auto
+    enable_block_reuse: false
+    free_gpu_memory_fraction: 0.85
+  moe_config:
+    backend: CUTLASS
+  print_iter_log: true

src/frameworks/a4x/trtllm-configs/deepseek-r1-nvfp4.yaml

@@ -182,6 +182,9 @@ spec:
               value: "{{ $root.Values.workload.model.name }}"
             - name: MODEL_DOWNLOAD_DIR
               value: "/ssd/{{ $root.Values.workload.model.name }}"
+            # A3-Ultra recipe is based on the TensorRT image, which puts tensorrt_llm in a different path than default
+            - name: TRTLLM_DIR
+              value: "/workspace/tensorrtllm_backend/tensorrt_llm"
             {{- if $root.Values.workload.envs }}
             {{- toYaml .Values.workload.envs | nindent 12 }}
             {{- end }}

src/helm-charts/a3ultra/inference-templates/deployment/templates/serving-launcher.yaml

@@ -0,0 +1,20 @@
+# Copyright 2025 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: single-host-serving-deployment-template
+description: single-host-serving-deployment-template
+type: application
+version: 0.1.0
+appVersion: "1.16.0"

src/helm-charts/a4x/inference-templates/deployment/Chart.yaml

@@ -0,0 +1,25 @@
+# Copyright 2025 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: ConfigMap
+metadata:
+  name: "{{ .Release.Name }}-config"
+data:
+  serving-configuration: |-
+{{- if .Values.serving_config  }}
+{{ .Values.serving_config | nindent 4 }}
+{{- else }}
+{{ "config: null" | nindent 4 }}
+{{- end }}

src/helm-charts/a4x/inference-templates/deployment/templates/serving-config-configmap.yaml

@@ -0,0 +1,27 @@
+# Copyright 2025 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: 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 }}