You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
cuda.core: map ml_dtypes bfloat16 to kDLBfloat in DLPack export - #3359
_smv_dtype_numpy_to_dlpack picked the DLPack type code from numpy.dtype.kind, but ml_dtypes.bfloat16 has kind "V", so a bf16 view fell through to the unsupported-dtype BufferError and could not be exported — even though the import side already accepts kDLBfloat (mapping it back to bfloat16 when ml_dtypes is installed) and _tensor_map.pyx maps bfloat16 the same way. A bf16 tensor from PyTorch/CuPy/JAX can be viewed but not re-exported.
This adds an explicit bfloat16 branch ahead of the kind chain that emits kDLBfloat with 16 bits, guarded by the module-level ml_dtypes availability check already used on the import side.
The new test round-trips through StridedMemoryView.from_dlpack itself, so it cannot skip on stacks where NumPy's own DLPack support for extended dtypes is unavailable — the hole that let the existing test_dlpack_export_roundtrip_dtypes[bfloat16] miss this via its NumPy-consumer probe.
Evidence from the released cuda-core 1.2.2 wheel (host memory; numpy 2.3.5, ml_dtypes 0.6.0), matching the issue's reproducer:
>>> view.__dlpack__()
BufferError: Unsupported dtype for DLPack export: bfloat16
with float16/uint16 views of the same buffer exporting fine. Additionally: cython -3 on cuda/core/_memoryview.pyx produces identical diagnostics on base and on this branch (the two existing cuStreamWaitEvent assign notes are pre-existing; none added), and the new branch's dtype equality check matches bfloat16 while not matching bool/uint16/float16/int8. Full build and suite validation runs in this repository's CI.
The documentation is up to date with these changes.
AI assistance disclosure: the change and test were prepared with an AI coding agent (GLM-5.3) under my direction; I reviewed each line, verified the reproducer and evidence above, and stand behind the change. The new test carries @pytest.mark.agent_authored(model="glm-5.3") per the provenance rules in AGENTS.md.
No actionable comments were generated in the recent review.
Summary
Summary by CodeRabbit
Bug Fixes
DLPack exports now correctly identify bfloat16 arrays, improving compatibility when sharing them with other tools. A round-trip check confirms that importing a bfloat16 array preserves its dtype, shape, and data pointer. Other data type handling remains unchanged. The bfloat16 check is skipped when the optional type support is unavailable.
DLPack export now maps ml_dtypes.bfloat16 to kDLBfloat. A round-trip test checks the imported dtype, shape, and data pointer.
Changes: bfloat16 DLPack export
Layer / File(s)
Summary
Map and round-trip bfloat16 cuda_core/cuda/core/_memoryview.pyx, cuda_core/tests/test_utils_dlpack.py
_smv_dtype_numpy_to_dlpack maps ml_dtypes.bfloat16 to kDLBfloat when ml_dtypes is available. The test checks the imported dtype, shape, and data pointer. It skips when ml_dtypes is unavailable.
Assessment against linked issues
Objective
Addressed
Explanation
Enable StridedMemoryView bfloat16 DLPack export and same-memory re-import [#3337]
✅
Suggested reviewers:rluo8
Change: Bug fix · Severity of issue fixed: Low
Merge Risk: Low · up to a8768
The export mapping is narrowly scoped and the test initializes CUDA correctly. Merge risk is low, with the remaining concern limited to skipped regression coverage in an incomplete test environment.
- [ ] Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Tip
🎯 Customize your risk assessments
Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.
✅ Automatic approvals, with your rules
Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.
Reviewing files that changed from the base of the PR and between 3d46b39 and a876815.
Files selected for processing: 1
cuda_core/tests/test_utils_dlpack.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
Reviewing files that changed from the base of the PR and between 4010017 and eb3fdb0.
Files selected for processing: 2
cuda_core/cuda/core/_memoryview.pyx
cuda_core/tests/test_utils_dlpack.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
Addressed: test_dlpack_export_bfloat16_roundtrip now requests init_cuda, since .view() reaches the driver via cuPointerGetAttributes — thanks for catching that this also matches the other driver-touching tests in this file.
Reviewing files that changed from the base of the PR and between eb3fdb0 and 3d46b39.
Files selected for processing: 1
cuda_core/tests/test_utils_dlpack.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Thanks, LGTM in general, but I'd ask you to move the branch and leave a bit more concrete bread-crumb that this test really should be irrelevant soon (and why).
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).
_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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
cuda.coreEverything related to the cuda.core module
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
seberg
left a comment
Description
closes #3337
_smv_dtype_numpy_to_dlpackpicked the DLPack type code fromnumpy.dtype.kind, butml_dtypes.bfloat16has kind"V", so a bf16 view fell through to the unsupported-dtypeBufferErrorand could not be exported — even though the import side already acceptskDLBfloat(mapping it back tobfloat16whenml_dtypesis installed) and_tensor_map.pyxmapsbfloat16the same way. A bf16 tensor from PyTorch/CuPy/JAX can be viewed but not re-exported.This adds an explicit bfloat16 branch ahead of the kind chain that emits
kDLBfloatwith 16 bits, guarded by the module-levelml_dtypesavailability check already used on the import side.The new test round-trips through
StridedMemoryView.from_dlpackitself, so it cannot skip on stacks where NumPy's own DLPack support for extended dtypes is unavailable — the hole that let the existingtest_dlpack_export_roundtrip_dtypes[bfloat16]miss this via its NumPy-consumer probe.Evidence from the released cuda-core 1.2.2 wheel (host memory; numpy 2.3.5, ml_dtypes 0.6.0), matching the issue's reproducer:
with float16/uint16 views of the same buffer exporting fine. Additionally:
cython -3oncuda/core/_memoryview.pyxproduces identical diagnostics on base and on this branch (the two existing cuStreamWaitEvent assign notes are pre-existing; none added), and the new branch's dtype equality check matchesbfloat16while not matching bool/uint16/float16/int8. Full build and suite validation runs in this repository's CI.Checklist
AI assistance disclosure: the change and test were prepared with an AI coding agent (GLM-5.3) under my direction; I reviewed each line, verified the reproducer and evidence above, and stand behind the change. The new test carries
@pytest.mark.agent_authored(model="glm-5.3")per the provenance rules in AGENTS.md.