diff --git a/cuda_core/cuda/core/_memoryview.pyx b/cuda_core/cuda/core/_memoryview.pyx index 287a01a70bc..aa5511253dd 100644 --- a/cuda_core/cuda/core/_memoryview.pyx +++ b/cuda_core/cuda/core/_memoryview.pyx @@ -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): + # ml_dtypes bfloat16 has kind "V", so the kind chain below cannot reach it. + code = kDLBfloat + elif kind == "b": if bits != 8: raise BufferError(f"Unsupported bool dtype itemsize: {np_dtype.itemsize}") code = kDLBool diff --git a/cuda_core/tests/test_utils_dlpack.py b/cuda_core/tests/test_utils_dlpack.py index 4664b2da3f3..997902ffb43 100644 --- a/cuda_core/tests/test_utils_dlpack.py +++ b/cuda_core/tests/test_utils_dlpack.py @@ -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. + """ + if ml_dtypes is None: + pytest.skip("ml_dtypes is not installed") + 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")],