Repository navigation
Add BF16 support to CPU and CUDA Pad kernels - #34062
Merged
Merged
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>
Copilot
AI
changed the title
[WIP] Add bfloat16 kernel for Pad node
Add BF16 support to CPU and CUDA Pad kernels
Nov 10, 2026
Tianlei Wu (tianleiwu)
marked this pull request as ready for review
November 10, 2026 04:43
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The updated BF16 registrations are missing from generated kernel documentation, causing validation to fail, and the CUDA regression test can falsely pass through CPU fallback.
2 open issues
What changed in this PR
Adds BF16 support to ONNX Runtime’s CPU and CUDA Pad kernels, addressing the session-initialization failure reported in #33874 for BF16 graphs that pad rotary caches.
Changes:
- Enables CPU BF16 Pad for opsets 21–25 with 16-bit constant-value handling.
- Adds CUDA BF16 registrations and kernel instantiations, including the optimized rank-2 path.
- Adds CPU- and CUDA-targeted tests for 2D bottom padding across opsets 21–25.
| File | Description |
|---|---|
onnxruntime/test/providers/cpu/tensor/pad_test.cc |
Adds BF16 padding regression tests for CPU and CUDA. |
onnxruntime/core/providers/cuda/tensor/pad.cc |
Instantiates BF16 Pad computation. |
onnxruntime/core/providers/cuda/tensor/pad_impl.cu |
Instantiates BF16 device kernels, including optimized padding. |
onnxruntime/core/providers/cuda/cuda_execution_provider.cc |
Registers BF16 Pad kernels with the CUDA provider. |
onnxruntime/core/providers/cpu/tensor/pad.cc |
Enables BF16 types and 16-bit value and computation dispatch. |
Review effort: Balanced
Add tensor(bfloat16) to the CPU Pad rows for opsets 21+ and to all CUDA Pad rows, matching the kernels registered by this PR. Generated with tools/python/gen_opkernel_doc.py (CPU and CUDA providers) on a local CUDA build and merged into the existing file; the DML and WebGPU sections are unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
RunCudaOnlyOnnxOpsetPadTest registered only the CUDA EP, but the CPU EP is always added implicitly, so the test could pass through CPU fallback when the CUDA kernel was missing. That became possible for BF16 now that the CPU EP also implements it. Set kOrtSessionOptionsDisableCPUEPFallback so a missing CUDA kernel fails session initialization instead. Co-authored-by: Copilot <223556219+Copilot@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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Description
Motivation and Context
BF16 Pad currently fails session initialization with CPU and CUDA providers, and CPU-only sessions report no implementation. Supporting it unblocks BF16 graphs that pad rotary caches without requiring a Concat workaround.