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.
DLPack exports now correctly recognize bfloat16, improving compatibility when sharing these arrays with other tools. Importing and exporting bfloat16 arrays preserves their dtype, shape, and data pointer. Handling of other data types remains unchanged.
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: Medium
Merge Risk: Low · up to 3d46b
The change is low risk, but the test should fail loudly when its declared ml-dtypes dependency is unavailable.
- [ ] 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 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.
The reason will be displayed to describe this comment to others. Learn more.
Major · Stability & Availability · Quick win
important: Establish a CUDA context for this test.
.view() queries Buffer.device_id, which calls cuPointerGetAttributes in cuda/core/_memory/_buffer.pyx. Without init_cuda, the test can depend on a context left by another test. (raw.githubusercontent.com)
Request init_cuda. As per coding guidelines: “Request init_cuda for any test that calls into the driver.”
_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>
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.
Minor · Functional Correctness · Quick win
important: Import ml_dtypes directly.
ml-dtypes is declared in both CUDA test dependency groups. (raw.githubusercontent.com) This skip hides a broken test environment and can leave the new export path untested. Remove the skip and use a direct import.
As per coding guidelines: “If the target is listed, import it directly.”
Source: Coding guidelines
This branch has not been deployed
No deployments
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
1 participant
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}
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.