rfcs : Added Embedding Bag RFC#5643
Conversation
- Added embedding bag RFC in rfcs/20260715-embedding-bag. Signed-off-by: Sudarshan Shinde <Sudarshan.Shinde@amd.com>
github-actions
Bot
added
documentation
vpirogov
left a comment
vpirogov
left a comment
There was a problem hiding this comment.
Thanks for the proposal, @sshinde-amd. Posted a few comments.
+@EikanWang for framework integration perspective.
| Embedding lookups and bag reductions are the dominant CPU operation in | ||
| recommendation models (DLRM, DCNv2, MMoE) and token-id-consuming NLP models. | ||
| PyTorch `nn.EmbeddingBag`, TensorFlow `embedding_lookup_sparse`, and ONNX | ||
| Runtime `EmbedLayerNormalization` each implement this on CPU today — either | ||
| through per-framework hand-rolled kernels or FBGEMM — because oneDNN offers no | ||
| equivalent primitive. |
There was a problem hiding this comment.
It would be great to include performance data to support this statement.
| Embedding Bag Operator on AMD CPUs] (https://dl.acm.org/doi/10.1109/MM.2024.3423785) ) | ||
| that embedding bag bandwidth can be improved by 9x over state of the art | ||
| implementations using different threading strategies.** |
There was a problem hiding this comment.
The paper does not show improvement over PyTorch baseline. Is there indication that Pytorch implementation is suboptimal?
| The first PR promotes the primitive plumbing to a public C++ API, adds the CPU | ||
| reference implementation, and validates the behavior needed by recommendation |
There was a problem hiding this comment.
New functionality is normally introduced as opt-in experimental feature first to avoid breaking library API during early adoption stages.
See Experimental Features section of the developer guide.
| PyTorch `nn.EmbeddingBag`, TensorFlow `embedding_lookup_sparse`, and ONNX | ||
| Runtime `EmbedLayerNormalization` each implement this on CPU today — either |
There was a problem hiding this comment.
Comparison of the semantics of these operations with semantics of proposed API would be very useful.
| oneDNN primitive uses Xbyak JIT. The embedding kernel logic (gather a row, | ||
| multiply by a scalar weight, reduce element-wise) is simple enough that there | ||
| is no shape-dependent code-gen opportunity that benefits from runtime | ||
| emission. ZenDNN's intrinsic implementation is well-validated today. The |
There was a problem hiding this comment.
There are two other reasons why oneDNN largely relies on Xbyak across platforms:
- Intrinsics require the compiler to support specific instruction set targeted. This would push minimal required compiler version in particular for newer instruction sets.
- Intrinsics were not compatible between GCC/Clang and MSVS compilers. Is this not the case anymore?
| 7. All MDs must be fully described. | ||
| 8. `dst_desc.dims[1] == src_desc.dims[1]` (embedding dim matches). | ||
|
|
||
| ## 4. PoC — embedding_bag primitive end-to-end with benchdnn |
There was a problem hiding this comment.
What does "end-to-end" mean here?
| `lowoha::embag` API. This RFC upstreams those intrinsic kernels into oneDNN as a | ||
| first-class primitive lifted into `src/cpu/x64/`, with ZenDNN removed from the | ||
| runtime path after integration. |
There was a problem hiding this comment.
The motivation part only covers reference implementation. Why this paragraph discusses intrinsics and ZenDNN?
| 2. **Threading strategy for very long bags.** Phase 1 parallelizes across bags | ||
| (`parallel_nd(B, ...)`). For workloads where a single bag is very long | ||
| relative to thread count, an optional inner split via `balance211` could | ||
| improve scaling. Deferred to Phase 2. |
There was a problem hiding this comment.
This does not seem relevant to the proposal.
2 participants
Description
This RFC adds an Embedding Bag as a public oneDNN primitive dnnl::embedding_bag, backed first by a portable CPU reference implementation. A thin dnnl::embedding wrapper covers the lookup-only (no-reduction) mode.
This primitive closely resembles PyTorch Embedding Bag semantics.
oneDNN currently has no primitive equivalent to PyTorch Embedding Bag. The closest dense neighbous (reduction, graph gather) do not cover indexed lookup + variable-length bag reduction semantics. Framework users that integrate with oneDNN either need to depend on its framework implementation, or assemble multiple ops by hand.
The first PR promotes the primitive plumbing to a public C++ API, adds the CPU reference implementation, and validates the behavior needed by recommendation and NLP workloads: sum / mean / max reduction over variable-length bags, optional per-sample weights, padding_idx skipping, and the include_last_offset flag matching PyTorch semantics. Future optimized implementations (AVX-512 intrinsics) reuse the same API without any framework changes.
https://github.com/amd/oneDNN-Zen/tree/zen-rfcs/rfcs/20260715-embedding-bag
Fixes # (github issue)
Checklist
General
make testandmake test_benchdnn_*) pass locally for each commit?Performance improvements
New features
Bug fixes
RFC PR