Skip to content

Commit 126fd7e

Browse files
yinggehyinggeh
committed
Add BF16 tests for ORT backend
1 parent 0844c37 commit 126fd7e

4 files changed

Lines changed: 281 additions & 0 deletions

File tree

‎qa/L0_backend_onnxruntime/bfloat16_test.py‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
#!/usr/bin/env python
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 sys
30+
import unittest
31+
32+
import ml_dtypes
33+
import numpy as np
34+
import tritonclient.grpc as grpcclient
35+
import tritonclient.http as httpclient
36+
37+
# Client type can be passed as first arg (e.g. python bfloat16_test.py http) or via CLIENT_TYPE env.
38+
if len(sys.argv) >= 2 and sys.argv[1] in ("http", "grpc"):
39+
os.environ["CLIENT_TYPE"] = sys.argv[1]
40+
del sys.argv[1]
41+
42+
43+
class BFloat16Test(unittest.TestCase):
44+
def setUp(self):
45+
self.protocol = os.environ.get("CLIENT_TYPE", "http")
46+
if self.protocol == "http":
47+
self.client_ = httpclient.InferenceServerClient("localhost:8000")
48+
else:
49+
self.client_ = grpcclient.InferenceServerClient("localhost:8001")
50+
self.model_name_ = "add_bf16"
51+
52+
def _assert_allclose_bf16(self, actual, desired, **kwargs):
53+
"""Compare bfloat16 arrays by converting to float32 for the check.
54+
55+
We cannot use np.testing.assert_allclose(actual, desired) directly:
56+
isclose() does result_type(y, 1.) and raises DTypePromotionError for
57+
bfloat16 in NumPy 2. The error message is misleading—it says
58+
"Float16DType and bfloat16" even when both arrays are bfloat16; the
59+
real conflict is bfloat16 vs the scalar 1.0 (float64) used inside
60+
isclose. Converting to float32 only for the comparison avoids this.
61+
"""
62+
np.testing.assert_allclose(
63+
np.asarray(actual, dtype=np.float32),
64+
np.asarray(desired, dtype=np.float32),
65+
**kwargs,
66+
)
67+
68+
def _infer_bf16(self, input0_data, input1_data):
69+
"""Helper to run BF16 inference and return the output numpy array."""
70+
if self.protocol == "http":
71+
input0 = httpclient.InferInput("INPUT0", [5, 5], "BF16")
72+
input1 = httpclient.InferInput("INPUT1", [5, 5], "BF16")
73+
else:
74+
input0 = grpcclient.InferInput("INPUT0", [5, 5], "BF16")
75+
input1 = grpcclient.InferInput("INPUT1", [5, 5], "BF16")
76+
input0.set_data_from_numpy(input0_data)
77+
input1.set_data_from_numpy(input1_data)
78+
79+
results = self.client_.infer(self.model_name_, [input0, input1])
80+
return results.as_numpy("OUTPUT")
81+
82+
def test_bf16_add_variants(self):
83+
"""Run multiple BF16 add cases in one test: zeros, negatives, large, small, cancellation, identical."""
84+
shape = (5, 5)
85+
86+
# Zeros: 0.0 + 0.0 = 0.0
87+
output = self._infer_bf16(
88+
np.zeros(shape, dtype=ml_dtypes.bfloat16),
89+
np.zeros(shape, dtype=ml_dtypes.bfloat16),
90+
)
91+
self.assertEqual(output.dtype, ml_dtypes.bfloat16)
92+
self._assert_allclose_bf16(output, np.zeros(shape, dtype=ml_dtypes.bfloat16))
93+
94+
# Negative and mixed: -1.5 + 3.5 = 2.0
95+
output = self._infer_bf16(
96+
np.full(shape, -1.5, dtype=ml_dtypes.bfloat16),
97+
np.full(shape, 3.5, dtype=ml_dtypes.bfloat16),
98+
)
99+
self.assertEqual(output.dtype, ml_dtypes.bfloat16)
100+
self._assert_allclose_bf16(
101+
output, np.full(shape, 2.0, dtype=ml_dtypes.bfloat16)
102+
)
103+
104+
# Large values within BF16 range: 100 + 200 = 300
105+
output = self._infer_bf16(
106+
np.full(shape, 100.0, dtype=ml_dtypes.bfloat16),
107+
np.full(shape, 200.0, dtype=ml_dtypes.bfloat16),
108+
)
109+
self.assertEqual(output.dtype, ml_dtypes.bfloat16)
110+
self._assert_allclose_bf16(
111+
output, np.full(shape, 300.0, dtype=ml_dtypes.bfloat16)
112+
)
113+
114+
# Small values (near underflow / precision limit): 0.01 + 0.01 = 0.02
115+
output = self._infer_bf16(
116+
np.full(shape, 1e-2, dtype=ml_dtypes.bfloat16),
117+
np.full(shape, 1e-2, dtype=ml_dtypes.bfloat16),
118+
)
119+
self.assertEqual(output.dtype, ml_dtypes.bfloat16)
120+
self._assert_allclose_bf16(
121+
output, np.full(shape, 2e-2, dtype=ml_dtypes.bfloat16)
122+
)
123+
124+
# Exact cancellation: 1.0 + (-1.0) = 0.0
125+
output = self._infer_bf16(
126+
np.full(shape, 1.0, dtype=ml_dtypes.bfloat16),
127+
np.full(shape, -1.0, dtype=ml_dtypes.bfloat16),
128+
)
129+
self.assertEqual(output.dtype, ml_dtypes.bfloat16)
130+
self._assert_allclose_bf16(output, np.zeros(shape, dtype=ml_dtypes.bfloat16))
131+
132+
# Identical inputs: 2.0 + 2.0 = 4.0
133+
output = self._infer_bf16(
134+
np.full(shape, 2.0, dtype=ml_dtypes.bfloat16),
135+
np.full(shape, 2.0, dtype=ml_dtypes.bfloat16),
136+
)
137+
self.assertEqual(output.dtype, ml_dtypes.bfloat16)
138+
self._assert_allclose_bf16(
139+
output, np.full(shape, 4.0, dtype=ml_dtypes.bfloat16)
140+
)
141+
142+
143+
if __name__ == "__main__":
144+
unittest.main()

‎qa/L0_backend_onnxruntime/models/add_bf16/1/model.onnx‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
 triton:w
2+

3+
INPUT0
4+
INPUT1OUTPUT"Addbf16_addZ
5+
INPUT0
6+

7+

8+
Z
9+
INPUT1
10+

11+

12+
b
13+
OUTPUT
14+

15+

16+
B

‎qa/L0_backend_onnxruntime/models/add_bf16/config.pbtxt‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# Redistribution and use in source and binary forms, with or without
4+
# modification, are permitted provided that the following conditions
5+
# are met:
6+
# * Redistributions of source code must retain the above copyright
7+
# notice, this list of conditions and the following disclaimer.
8+
# * Redistributions in binary form must reproduce the above copyright
9+
# notice, this list of conditions and the following disclaimer in the
10+
# documentation and/or other materials provided with the distribution.
11+
# * Neither the name of NVIDIA CORPORATION nor the names of its
12+
# contributors may be used to endorse or promote products derived
13+
# from this software without specific prior written permission.
14+
#
15+
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
16+
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17+
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18+
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
19+
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20+
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21+
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22+
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
23+
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24+
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
25+
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26+
27+
platform: "onnxruntime_onnx"
28+
max_batch_size: 0
29+
input [
30+
{
31+
name: "INPUT0"
32+
data_type: TYPE_BF16
33+
dims: [5, 5]
34+
},
35+
{
36+
name: "INPUT1"
37+
data_type: TYPE_BF16
38+
dims: [5, 5]
39+
}
40+
]
41+
output [
42+
{
43+
name: "OUTPUT"
44+
data_type: TYPE_BF16
45+
dims: [5, 5]
46+
}
47+
]
48+
instance_group: {
49+
kind: KIND_GPU
50+
}

‎qa/L0_backend_onnxruntime/test.sh‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
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+
37+
# BFLOAT16 test
38+
SERVER_ARGS="--model-repository=`pwd`/models"
39+
run_server
40+
if [ "$SERVER_PID" == "0" ]; then
41+
echo -e "\n***\n*** Failed to start $SERVER\n***"
42+
cat $SERVER_LOG
43+
exit 1
44+
fi
45+
46+
RET=0
47+
48+
set +e
49+
50+
for client_type in http grpc; do
51+
CLIENT_LOG="./bfloat16_test_${client_type}.log"
52+
python bfloat16_test.py $client_type >>$CLIENT_LOG 2>&1
53+
if [ $? -ne 0 ]; then
54+
cat $CLIENT_LOG
55+
echo -e "\n***\n*** Test Failed ($client_type)\n***"
56+
RET=1
57+
fi
58+
done
59+
60+
set -e
61+
62+
kill $SERVER_PID
63+
wait $SERVER_PID
64+
65+
if [ $RET -eq 0 ]; then
66+
echo -e "\n***\n*** Test Passed\n***"
67+
else
68+
echo -e "\n***\n*** Test FAILED\n***"
69+
fi
70+
71+
exit $RET

0 commit comments

Comments
 (0)