Skip to content

Commit 3d46b39

Browse files
cuda.core: map ml_dtypes bfloat16 to kDLBfloat in DLPack export
_smv_dtype_numpy_to_dlpack dispatched on numpy dtype kind, but ml_dtypes.bfloat16 has kind "V", so bf16 views fell through to the unsupported-dtype BufferError and could not be re-exported, even though the DLPack import path already accepts kDLBfloat and _tensor_map.pyx maps bfloat16. Detect bfloat16 up front and export it as kDLBfloat (16 bits). The new round-trip test goes through StridedMemoryView.from_dlpack so it cannot skip on stacks where NumPy's own DLPack support for extended dtypes is unavailable, unlike the existing parametrized round-trip test's NumPy-consumer probe. Signed-off-by: Sainath-Reddy7 <180950590+Sainath-Reddy7@users.noreply.github.com>
1 parent 4010017 commit 3d46b39

2 files changed

Lines changed: 24 additions & 1 deletion

File tree

‎cuda_core/cuda/core/_memoryview.pyx‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -669,7 +669,10 @@ cdef inline int _smv_dtype_numpy_to_dlpack(object dtype_obj, DLDataType* out_dty
669669
cdef str kind = np_dtype.kind
670670
cdef int bits = np_dtype.itemsize * 8
671671
cdef uint8_t code
672-
if kind == "b":
672+
if bfloat16 is not None and np_dtype == numpy.dtype(bfloat16):
673+
# ml_dtypes bfloat16 has kind "V", so the kind chain below cannot reach it.
674+
code = <uint8_t>kDLBfloat
675+
elif kind == "b":
673676
if bits != 8:
674677
raise BufferError(f"Unsupported bool dtype itemsize: {np_dtype.itemsize}")
675678
code = <uint8_t>kDLBool

‎cuda_core/tests/test_utils_dlpack.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,26 @@ def test_dlpack_export_roundtrip_dtypes(dtype):
7171
_assert_dlpack_export_roundtrip(np.zeros((2, 3), dtype=dtype))
7272

7373

74+
@pytest.mark.agent_authored(model="glm-5.3")
75+
def test_dlpack_export_bfloat16_roundtrip(init_cuda):
76+
"""bfloat16 views export as kDLBfloat and re-import as bfloat16 (#3337).
77+
78+
Unlike ``test_dlpack_export_roundtrip_dtypes[bfloat16]``, this round-trips
79+
through ``StridedMemoryView.from_dlpack`` itself, so it cannot skip on
80+
stacks where NumPy's own DLPack support for extended dtypes is unavailable.
81+
"""
82+
if ml_dtypes is None:
83+
pytest.skip("ml_dtypes is not installed")
84+
src = np.zeros((2, 3), dtype=np.uint16)
85+
view = StridedMemoryView.from_any_interface(src, stream_ptr=-1).view(
86+
dtype=np.dtype(ml_dtypes.bfloat16)
87+
)
88+
imported = StridedMemoryView.from_dlpack(view, stream_ptr=-1)
89+
assert imported.dtype == np.dtype(ml_dtypes.bfloat16)
90+
assert imported.shape == (2, 3)
91+
assert imported.ptr == view.ptr
92+
93+
7494
@pytest.mark.parametrize(
7595
"shape",
7696
[pytest.param((), id="scalar"), pytest.param((0, 3), id="empty")],

0 commit comments

Comments
 (0)