Skip to content

Commit d159373

Browse files
pskiran1pskiran1
authored
Merge branch 'main' into spolisetty/tri-669-psirt-triton-inference-server-vertex-ai-integration-access
2 parents a8da693 + 952a9fb commit d159373

4 files changed

Lines changed: 279 additions & 9 deletions

File tree

‎qa/L0_backend_onnxruntime/gen_add_bf16_onnx_model.py‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
#!/usr/bin/env python3
2+
# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
#
4+
# Redistribution and use in source and binary forms, with or without
5+
# modification, are permitted provided that the following conditions
6+
# are met:
7+
# * Redistributions of source code must retain the above copyright
8+
# notice, this list of conditions and the following disclaimer.
9+
# * Redistributions in binary form must reproduce the above copyright
10+
# notice, this list of conditions and the following disclaimer in the
11+
# documentation and/or other materials provided with the distribution.
12+
# * Neither the name of NVIDIA CORPORATION nor the names of its
13+
# contributors may be used to endorse or promote products derived
14+
# from this software without specific prior written permission.
15+
#
16+
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
17+
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18+
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
19+
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
20+
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
21+
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
22+
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
23+
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
24+
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
25+
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
26+
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
27+
28+
29+
# Generates the add_bf16 ONNX model and Triton config.
30+
# Model: element-wise Add in BFLOAT16 (INPUT0 + INPUT1 = OUTPUT), ONNX Runtime backend.
31+
import os
32+
33+
import onnx
34+
35+
36+
def generate_bf16_add_model(models_dir):
37+
"""Generate a simple BFLOAT16 Add model (INPUT0 + INPUT1 = OUTPUT)."""
38+
model_name = "add_bf16"
39+
shape = [1]
40+
onnx_dtype = onnx.TensorProto.BFLOAT16
41+
42+
add = onnx.helper.make_node("Add", ["INPUT0", "INPUT1"], ["OUTPUT"])
43+
44+
input0 = onnx.helper.make_tensor_value_info("INPUT0", onnx_dtype, shape)
45+
input1 = onnx.helper.make_tensor_value_info("INPUT1", onnx_dtype, shape)
46+
output = onnx.helper.make_tensor_value_info("OUTPUT", onnx_dtype, shape)
47+
48+
graph_proto = onnx.helper.make_graph(
49+
[add],
50+
model_name,
51+
[input0, input1],
52+
[output],
53+
)
54+
model_def = onnx.helper.make_model(graph_proto, producer_name="triton")
55+
# Cap IR version for older ONNX Runtime (e.g. max supported 11)
56+
model_def.ir_version = min(model_def.ir_version, 11)
57+
# BFLOAT16 support requires opset 13+
58+
model_def.opset_import[0].version = 13
59+
60+
model_dir = os.path.join(models_dir, model_name, "1")
61+
os.makedirs(model_dir, exist_ok=True)
62+
onnx.save(model_def, os.path.join(model_dir, "model.onnx"))
63+
64+
# Write config.pbtxt
65+
config = """platform: "onnxruntime_onnx"
66+
max_batch_size: 0
67+
input [
68+
{{
69+
name: "INPUT0"
70+
data_type: TYPE_BF16
71+
dims: {shape}
72+
}},
73+
{{
74+
name: "INPUT1"
75+
data_type: TYPE_BF16
76+
dims: {shape}
77+
}}
78+
]
79+
output [
80+
{{
81+
name: "OUTPUT"
82+
data_type: TYPE_BF16
83+
dims: {shape}
84+
}}
85+
]
86+
""".format(
87+
shape=shape
88+
)
89+
90+
config_path = os.path.join(models_dir, model_name, "config.pbtxt")
91+
with open(config_path, "w") as f:
92+
f.write(config)
93+
94+
print(f"Generated model '{model_name}' in {models_dir}")
95+
96+
97+
if __name__ == "__main__":
98+
models_dir = os.path.join(os.getcwd(), "models")
99+
os.makedirs(models_dir, exist_ok=True)
100+
generate_bf16_add_model(models_dir)

‎qa/L0_backend_onnxruntime/test.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
#!/usr/bin/env python3
2+
# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
#
4+
# Redistribution and use in source and binary forms, with or without
5+
# modification, are permitted provided that the following conditions
6+
# are met:
7+
# * Redistributions of source code must retain the above copyright
8+
# notice, this list of conditions and the following disclaimer.
9+
# * Redistributions in binary form must reproduce the above copyright
10+
# notice, this list of conditions and the following disclaimer in the
11+
# documentation and/or other materials provided with the distribution.
12+
# * Neither the name of NVIDIA CORPORATION nor the names of its
13+
# contributors may be used to endorse or promote products derived
14+
# from this software without specific prior written permission.
15+
#
16+
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
17+
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18+
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
19+
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
20+
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
21+
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
22+
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
23+
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
24+
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
25+
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
26+
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
27+
28+
import os
29+
import unittest
30+
31+
import numpy as np
32+
import tritonclient.grpc as grpcclient
33+
import tritonclient.http as httpclient
34+
35+
36+
class BFloat16Test(unittest.TestCase):
37+
def setUp(self):
38+
self.protocol = os.environ.get("CLIENT_TYPE", "http")
39+
if self.protocol == "http":
40+
self.client_ = httpclient.InferenceServerClient("localhost:8000")
41+
else:
42+
self.client_ = grpcclient.InferenceServerClient("localhost:8001")
43+
self.model_name_ = "add_bf16"
44+
self.shape_ = [1]
45+
46+
def _infer_bf16(self, input0_data, input1_data):
47+
"""Helper to run BF16 inference and return the output numpy array."""
48+
if self.protocol == "http":
49+
input0 = httpclient.InferInput("INPUT0", self.shape_, "BF16")
50+
input1 = httpclient.InferInput("INPUT1", self.shape_, "BF16")
51+
else:
52+
input0 = grpcclient.InferInput("INPUT0", self.shape_, "BF16")
53+
input1 = grpcclient.InferInput("INPUT1", self.shape_, "BF16")
54+
input0.set_data_from_numpy(input0_data)
55+
input1.set_data_from_numpy(input1_data)
56+
57+
results = self.client_.infer(self.model_name_, [input0, input1])
58+
return results.as_numpy("OUTPUT")
59+
60+
def test_bf16_add_variants(self):
61+
"""Run BF16 add across multiple cases: zeros, negatives, large, small, cancellation, and identical."""
62+
for input0_val, input1_val, expected_val in [
63+
(0.0, 0.0, 0.0), # zeros
64+
(-1.5, 3.5, 2.0), # negatives / mixed
65+
(100.0, 200.0, 300.0), # large
66+
(1e-2, 1e-2, 2e-2), # small (near underflow)
67+
(1.0, -1.0, 0.0), # cancellation
68+
(2.0, 2.0, 4.0), # identical inputs
69+
]:
70+
output = self._infer_bf16(
71+
np.full(self.shape_, input0_val, dtype=np.float32),
72+
np.full(self.shape_, input1_val, dtype=np.float32),
73+
)
74+
self.assertEqual(output.dtype, np.float32)
75+
# TODO: BF16 to FP32 conversion loses precision. Remove rtol and atol in TRI-801.
76+
# BF16 has ~3 decimal digits; use relaxed tol for computed values
77+
np.testing.assert_allclose(output, expected_val, rtol=1e-2, atol=1e-3)
78+
79+
80+
if __name__ == "__main__":
81+
unittest.main()

‎qa/L0_backend_onnxruntime/test.sh‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
#!/bin/bash
2+
# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
#
4+
# Redistribution and use in source and binary forms, with or without
5+
# modification, are permitted provided that the following conditions
6+
# are met:
7+
# * Redistributions of source code must retain the above copyright
8+
# notice, this list of conditions and the following disclaimer.
9+
# * Redistributions in binary form must reproduce the above copyright
10+
# notice, this list of conditions and the following disclaimer in the
11+
# documentation and/or other materials provided with the distribution.
12+
# * Neither the name of NVIDIA CORPORATION nor the names of its
13+
# contributors may be used to endorse or promote products derived
14+
# from this software without specific prior written permission.
15+
#
16+
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
17+
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18+
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
19+
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
20+
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
21+
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
22+
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
23+
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
24+
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
25+
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
26+
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
27+
28+
export CUDA_VISIBLE_DEVICES=0
29+
30+
SERVER=/opt/tritonserver/bin/tritonserver
31+
SERVER_LOG="./inference_server.log"
32+
CLIENT_LOG="./test.log"
33+
source ../common/util.sh
34+
35+
rm -f *.log
36+
rm -rf models
37+
38+
RET=0
39+
40+
# BFLOAT16 test
41+
# Generate the model
42+
mkdir -p models/add_bf16/1
43+
set +e
44+
45+
pip install onnx==1.20.1
46+
if [ $? -ne 0 ]; then
47+
echo -e "\n***\n*** Failed to install onnx dependency\n***"
48+
exit 1
49+
fi
50+
51+
python gen_add_bf16_onnx_model.py
52+
if [ $? -ne 0 ]; then
53+
echo -e "\n***\n*** Failed to generate BFLOAT16 ONNX model\n***"
54+
exit 1
55+
fi
56+
57+
set -e
58+
59+
SERVER_ARGS="--model-repository=`pwd`/models"
60+
run_server
61+
if [ "$SERVER_PID" == "0" ]; then
62+
echo -e "\n***\n*** Failed to start $SERVER\n***"
63+
cat $SERVER_LOG
64+
exit 1
65+
fi
66+
67+
set +e
68+
69+
for client_type in http grpc; do
70+
export CLIENT_TYPE=$client_type
71+
CLIENT_LOG="./test_${client_type}.log"
72+
python test.py >>$CLIENT_LOG 2>&1
73+
if [ $? -ne 0 ]; then
74+
cat $CLIENT_LOG
75+
echo -e "\n***\n*** Test Failed ($client_type)\n***"
76+
RET=1
77+
fi
78+
done
79+
unset CLIENT_TYPE
80+
81+
set -e
82+
83+
kill $SERVER_PID
84+
wait $SERVER_PID
85+
86+
if [ $RET -eq 0 ]; then
87+
echo -e "\n***\n*** Test Passed\n***"
88+
else
89+
echo -e "\n***\n*** Test FAILED\n***"
90+
fi
91+
92+
exit $RET

‎qa/common/test_util.py‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#!/usr/bin/env python3
22

3-
# Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
# Copyright 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
44
#
55
# Redistribution and use in source and binary forms, with or without
66
# modification, are permitted provided that the following conditions
@@ -313,16 +313,13 @@ def check_gpus_compute_capability(min_capability):
313313
import importlib.util
314314

315315
if importlib.util.find_spec("cuda") is not None:
316-
import cuda.core.experimental as cuda_core_experimental
317-
318-
devices = cuda_core_experimental.system.devices
316+
from cuda.core import Device
319317

318+
devices = Device.get_all_devices()
320319
for device in devices:
321-
compute_capability = (
322-
device.compute_capability.major + device.compute_capability.minor / 10.0
323-
)
324-
325-
if compute_capability < min_capability:
320+
cc = device.compute_capability
321+
compute_capability_value = cc.major + cc.minor / 10.0
322+
if compute_capability_value < min_capability:
326323
return False
327324

328325
elif importlib.util.find_spec("pycuda") is not None:

0 commit comments

Comments
 (0)