Skip to content

Commit 6b8430a

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 <[email protected]>
1 parent 4010017 commit 6b8430a

2 files changed

Lines changed: 21 additions & 0 deletions

File tree

‎cuda_core/cuda/core/_memoryview.pyx‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -689,6 +689,10 @@ cdef inline int _smv_dtype_numpy_to_dlpack(object dtype_obj, DLDataType* out_dty
689689
if bits not in (64, 128):
690690
raise BufferError(f"Unsupported complex dtype: {np_dtype}")
691691
code = <uint8_t>kDLComplex
692+
elif np_dtype.type is bfloat16:
693+
# ml_dtypes bfloat16 has kind "V", so the kind chain cannot reach it;
694+
# the byte-order check above has already run for it.
695+
code = <uint8_t>kDLBfloat
692696
else:
693697
raise BufferError(f"Unsupported dtype for DLPack export: {np_dtype}")
694698

‎cuda_core/tests/test_utils_dlpack.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,23 @@ 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+
This test becomes unnecessary when `_assert_dlpack_export_roundtrip`
79+
starts testing bfloat16, which should be `ml_dtypes>0.6.0`.
80+
"""
81+
if ml_dtypes is None:
82+
pytest.skip("ml_dtypes is not installed")
83+
src = np.zeros((2, 3), dtype=np.uint16)
84+
view = StridedMemoryView.from_any_interface(src, stream_ptr=-1).view(dtype=np.dtype(ml_dtypes.bfloat16))
85+
imported = StridedMemoryView.from_dlpack(view, stream_ptr=-1)
86+
assert imported.dtype == np.dtype(ml_dtypes.bfloat16)
87+
assert imported.shape == (2, 3)
88+
assert imported.ptr == view.ptr
89+
90+
7491
@pytest.mark.parametrize(
7592
"shape",
7693
[pytest.param((), id="scalar"), pytest.param((0, 3), id="empty")],

0 commit comments

Comments
 (0)