From a8768150166244e628bc97e4a3ff0815ce8b3060 Mon Sep 17 00:00:00 2001
From: Sainath-Reddy7 <180950590+Sainath-Reddy7@users.noreply.github.com>
Date: Wed, 4 Nov 2026 12:14:37 +0530
Subject: [PATCH] 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>
---
 cuda_core/cuda/core/_memoryview.pyx  |  5 ++++-
 cuda_core/tests/test_utils_dlpack.py | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

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 = <uint8_t>kDLBfloat
+    elif kind == "b":
         if bits != 8:
             raise BufferError(f"Unsupported bool dtype itemsize: {np_dtype.itemsize}")
         code = <uint8_t>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")],