Skip to content

feat: Enable bfloat16 input/output tensor dtype in Python client #880

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
yinggeh wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

feat: Enable bfloat16 input/output tensor dtype in Python client #880

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e39efb7
Remove special handling to bfloat16
yinggeh Feb 14, 2026
536e75a
Remove commented code initially for debugging purpose
yinggeh Feb 17, 2026
a1618ea
Remove unused functions
yinggeh Feb 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/python/library/requirements/requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Copyright 2020-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
Expand All @@ -24,6 +24,7 @@
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

ml_dtypes>=0.5.4
numpy>=1.19.1
python-rapidjson>=0.9.1
urllib3>=2.0.7
29 changes: 7 additions & 22 deletions src/python/library/tritonclient/grpc/_infer_input.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
#!/usr/bin/env python3

# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
Expand Down Expand Up @@ -124,22 +124,13 @@ def set_data_from_numpy(self, input_tensor):
"""
if not isinstance(input_tensor, (np.ndarray,)):
raise_error("input_tensor must be a numpy array")
# DLIS-3986: Special handling for bfloat16 until Numpy officially supports it
if self._input.datatype == "BF16":
if input_tensor.dtype != triton_to_np_dtype(self._input.datatype):
raise_error(
"got unexpected datatype {} from numpy array, expected {} for BF16 type".format(
input_tensor.dtype, triton_to_np_dtype(self._input.datatype)
)
)
else:
dtype = np_to_triton_dtype(input_tensor.dtype)
if self._input.datatype != dtype:
raise_error(
"got unexpected datatype {} from numpy array, expected {}".format(
dtype, self._input.datatype
)
dtype = np_to_triton_dtype(input_tensor.dtype)
if self._input.datatype != dtype:
raise_error(
"got unexpected datatype {} from numpy array, expected {}".format(
dtype, self._input.datatype
)
)
Comment on lines +127 to +133

Copilot AI Feb 17, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change makes BF16 inputs require an actual BF16 numpy dtype (via ml_dtypes.bfloat16 mapping). Previously, BF16 inputs could be provided as float32 (since triton_to_np_dtype("BF16") mapped to np.float32) and the client handled BF16 serialization. If existing users rely on the float32 workaround, consider accepting both float32 and bfloat16 for BF16 inputs (converting float32 to BF16 bytes) or clearly documenting the breaking behavior change.

Copilot uses AI. Check for mistakes.
valid_shape = True
if len(self._input.shape) != len(input_tensor.shape):
valid_shape = False
Expand All @@ -163,12 +154,6 @@ def set_data_from_numpy(self, input_tensor):
self._raw_content = serialized_output.item()
else:
self._raw_content = b""
elif self._input.datatype == "BF16":
serialized_output = serialize_bf16_tensor(input_tensor)
if serialized_output.size > 0:
self._raw_content = serialized_output.item()
else:
self._raw_content = b""
else:
self._raw_content = input_tensor.tobytes()
return self
Expand Down
6 changes: 1 addition & 5 deletions src/python/library/tritonclient/grpc/_infer_result.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
#!/usr/bin/env python3

# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
Expand Down Expand Up @@ -77,10 +77,6 @@ def as_numpy(self, name):
np_array = deserialize_bytes_tensor(
self._result.raw_output_contents[index]
)
elif datatype == "BF16":
np_array = deserialize_bf16_tensor(
self._result.raw_output_contents[index]
)
else:
np_array = np.frombuffer(
self._result.raw_output_contents[index],
Expand Down
38 changes: 16 additions & 22 deletions src/python/library/tritonclient/http/_infer_input.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
#!/usr/bin/env python3

# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
Expand Down Expand Up @@ -129,22 +129,22 @@
"""
if not isinstance(input_tensor, (np.ndarray,)):
raise_error("input_tensor must be a numpy array")
# DLIS-3986: Special handling for bfloat16 until Numpy officially supports it
if self._datatype == "BF16":
if input_tensor.dtype != triton_to_np_dtype(self._datatype):
raise_error(
"got unexpected datatype {} from numpy array, expected {} for BF16 type".format(
input_tensor.dtype, triton_to_np_dtype(self._datatype)
)
)
else:
dtype = np_to_triton_dtype(input_tensor.dtype)
if self._datatype != dtype:
raise_error(
"got unexpected datatype {} from numpy array, expected {}".format(
dtype, self._datatype
)
# # DLIS-3986: Special handling for bfloat16 until Numpy officially supports it
# if self._datatype == "BF16":
# if input_tensor.dtype != triton_to_np_dtype(self._datatype):
# raise_error(
# "got unexpected datatype {} from numpy array, expected {} for BF16 type".format(
# input_tensor.dtype, triton_to_np_dtype(self._datatype)
Comment thread
github-advanced-security[bot] marked this conversation as resolved.
Fixed
# )
# )
# else:
dtype = np_to_triton_dtype(input_tensor.dtype)
if self._datatype != dtype:
raise_error(
"got unexpected datatype {} from numpy array, expected {}".format(
dtype, self._datatype
)
)
Comment thread
yinggeh marked this conversation as resolved.
Outdated
Comment on lines +126 to +132

Copilot AI Feb 17, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change makes BF16 inputs require input_tensor.dtype to map to "BF16" (i.e., ml_dtypes.bfloat16). Previously the client accepted float32 tensors for BF16 (via triton_to_np_dtype("BF16") mapping) and handled conversion/serialization. If backward compatibility is desired, consider allowing float32 for BF16 here (and converting to BF16 bytes) or explicitly documenting this breaking change in the client behavior.

Copilot uses AI. Check for mistakes.
valid_shape = True
if len(self._shape) != len(input_tensor.shape):
valid_shape = False
Expand Down Expand Up @@ -202,12 +202,6 @@
self._raw_data = serialized_output.item()
else:
self._raw_data = b""
elif self._datatype == "BF16":
serialized_output = serialize_bf16_tensor(input_tensor)
if serialized_output.size > 0:
self._raw_data = serialized_output.item()
else:
self._raw_data = b""
else:
self._raw_data = input_tensor.tobytes()
Comment on lines -184 to 191

Copilot AI Feb 17, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With the BF16-specific serialization branch removed, this method now always uses input_tensor.tobytes() for non-BYTES dtypes. The file still imports serialize_bf16_tensor, which appears unused now; please remove that import to avoid unused/dead code paths.

Copilot uses AI. Check for mistakes.
self._parameters["binary_data_size"] = len(self._raw_data)
Expand Down
6 changes: 1 addition & 5 deletions src/python/library/tritonclient/http/_infer_result.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
#!/usr/bin/env python3

# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
Expand Down Expand Up @@ -190,10 +190,6 @@ def as_numpy(self, name):
np_array = deserialize_bytes_tensor(
self._buffer[start_index:end_index]
)
elif datatype == "BF16":
np_array = deserialize_bf16_tensor(
self._buffer[start_index:end_index]
)
else:
np_array = np.frombuffer(
self._buffer[start_index:end_index],
Expand Down
9 changes: 7 additions & 2 deletions src/python/library/tritonclient/utils/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
#!/usr/bin/env python3

# Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
Expand Down Expand Up @@ -28,6 +28,7 @@

import struct

import ml_dtypes
import numpy as np

from ._shared_memory_tensor import SharedMemoryTensor
Expand Down Expand Up @@ -149,6 +150,8 @@ def np_to_triton_dtype(np_dtype):
return "UINT32"
elif np_dtype == np.uint64:
return "UINT64"
elif np_dtype == ml_dtypes.bfloat16:
return "BF16"
elif np_dtype == np.float16:
return "FP16"
elif np_dtype == np.float32:
Expand Down Expand Up @@ -179,9 +182,11 @@ def triton_to_np_dtype(dtype):
return np.uint32
elif dtype == "UINT64":
return np.uint64
elif dtype == "BF16":
return ml_dtypes.bfloat16
elif dtype == "FP16":
return np.float16
elif dtype == "FP32" or dtype == "BF16":
elif dtype == "FP32":
return np.float32
elif dtype == "FP64":
return np.float64
Expand Down
Loading