Skip to content

cuda.core: map ml_dtypes bfloat16 to kDLBfloat in DLPack export #3359

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

Open
Sainath-Reddy7 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

cuda.core: map ml_dtypes bfloat16 to kDLBfloat in DLPack export #3359

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
5 changes: 4 additions & 1 deletion cuda_core/cuda/core/_memoryview.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -669,7 +669,10 @@ cdef inline int _smv_dtype_numpy_to_dlpack(object dtype_obj, DLDataType* out_dty
cdef str kind = np_dtype.kind
cdef int bits = np_dtype.itemsize * 8
cdef uint8_t code
if kind == "b":
if bfloat16 is not None and np_dtype == numpy.dtype(bfloat16):

@seberg seberg Nov 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Can you move this to be the last branch here, and also can you make the check np_dtype.type is bfloat16 (we already checked byte-order, so that is fine).

# ml_dtypes bfloat16 has kind "V", so the kind chain below cannot reach it.
code = <uint8_t>kDLBfloat
elif kind == "b":
if bits != 8:
raise BufferError(f"Unsupported bool dtype itemsize: {np_dtype.itemsize}")
code = <uint8_t>kDLBool
Expand Down
18 changes: 18 additions & 0 deletions cuda_core/tests/test_utils_dlpack.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,24 @@ def test_dlpack_export_roundtrip_dtypes(dtype):
_assert_dlpack_export_roundtrip(np.zeros((2, 3), dtype=dtype))


@pytest.mark.agent_authored(model="glm-5.3")
def test_dlpack_export_bfloat16_roundtrip(init_cuda):
"""bfloat16 views export as kDLBfloat and re-import as bfloat16 (#3337).

Unlike ``test_dlpack_export_roundtrip_dtypes[bfloat16]``, this round-trips
through ``StridedMemoryView.from_dlpack`` itself, so it cannot skip on
stacks where NumPy's own DLPack support for extended dtypes is unavailable.
Comment on lines +77 to +80

@seberg seberg Nov 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Suggested change
Unlike ``test_dlpack_export_roundtrip_dtypes[bfloat16]``, this round-trips
through ``StridedMemoryView.from_dlpack`` itself, so it cannot skip on
stacks where NumPy's own DLPack support for extended dtypes is unavailable.
This test becomes unnecessary when `_assert_dlpack_export_roundtrip`
starts testing bfloat16, which should be `ml_dtypes>0.6.0`.

Let's just add a bread-crumb that we don't really need this test. The above already would hit the issue if we installed a newer ml_dtypes

"""
if ml_dtypes is None:
pytest.skip("ml_dtypes is not installed")
Comment thread
seberg marked this conversation as resolved.
Resolved
src = np.zeros((2, 3), dtype=np.uint16)
view = StridedMemoryView.from_any_interface(src, stream_ptr=-1).view(dtype=np.dtype(ml_dtypes.bfloat16))
imported = StridedMemoryView.from_dlpack(view, stream_ptr=-1)
assert imported.dtype == np.dtype(ml_dtypes.bfloat16)
assert imported.shape == (2, 3)
assert imported.ptr == view.ptr


@pytest.mark.parametrize(
"shape",
[pytest.param((), id="scalar"), pytest.param((0, 3), id="empty")],
Expand Down
Loading