Repository navigation
fbgemm-xpu: add all_to_one_device for XPU - #140
Merged
Merged
Conversation
TorchRec's SeqEmbeddingsAllToOne gathers per-rank sequence embedding outputs onto the output device with fbgemm::all_to_one_device. Without an XPU implementation, sharded QuantEmbeddingCollection inference across XPU devices fails in output distribution with NotImplementedError. The XPU version keeps the CUDA contract: inputs must be XPU tensors, the target an XPU device with an index, and tensors already on the target are returned as they are. The others are copied with a non-blocking Tensor.to under a guard for the target device, so PyTorch's XPU cross-device copy orders the streams. The schema is defined only when FBGEMM has not registered it. The cross-device test needs two XPU devices and is skipped otherwise, like the two-device cases in test_jagged_device_validation.py. Signed-off-by: mkrze <mateusz.krzemieniewski@intel.com>
Mirror FBGEMM's test_all_to_one_device: inputs round-robin over every visible XPU, including row-strided views, gathered onto each device in turn. Inputs already on the target must come back as the same views. With one device the test still runs and checks the aliasing path. Signed-off-by: mkrze <mateusz.krzemieniewski@intel.com>
The Bandit Scanner workflow fails on B101 for bare asserts. Mark them with "# nosec B101" as the other plugin tests do. Signed-off-by: mkrze <mateusz.krzemieniewski@intel.com>
The package depends on torch~=2.14.0, so pip resolves it to the newest 2.14.x. Since 2.14.1+xpu showed up on the XPU index, the exact version check in the xpu fixture turns every test in the file into a setup error (122 errors on bmg). Compare the minor version, which is what the dependency actually promises. Signed-off-by: mkrze <mateusz.krzemieniewski@intel.com>
mkrze
changed the title
mkrze
marked this pull request as ready for review
October 6, 2026 08:14
mkrze
requested review from
aagalleg,
dvrogozh,
flezaalv and
molintc
as code owners
October 6, 2026 08:14
aagalleg
requested changes
Oct 6, 2026
aagalleg
left a comment
aagalleg
left a comment
Contributor
There was a problem hiding this comment.
Please address the following changes.
dvrogozh
reviewed
Oct 6, 2026
| # Copyright (c) 2026 Intel Corporation. All Rights Reserved. | ||
| # SPDX-License-Identifier: BSD-3-Clause | ||
|
|
||
| """``all_to_one_device`` on XPU. |
Contributor
There was a problem hiding this comment.
Are there any tests we can reuse from fbgemm?
Contributor
Author
There was a problem hiding this comment.
Yes, FBGEMM has a test for it.
Addressed in 993e8bb with a hybrid approach: we now run the upstream test on XPU through the shared patch and dropped our copy of it, and keep a few local tests for what upstream doesn't check (no copy for tensors already on the target, contiguous copies, invalid devices, non-default streams).
Contributor
There was a problem hiding this comment.
Thank you. One thing I really care about is reusing fbgemm tests as much as possible and keeping our own tests to the bare minimum. This way we increase our chances to be really compatible with the upstream and minimize our debt for potential upstreaming.
Without peer access between the source and target XPU, PyTorch's cross-device copy stages the data through pageable host memory. With non_blocking set, the device-to-host half is not waited on before the host reads the staging buffer, so the result can be stale. Use non_blocking only when at::xpu::get_p2p_access reports peer access for the pair, the same query the XPU copy uses to choose its path, and copy synchronously otherwise. CUDA instead requires peer access between all devices; the fallback keeps systems without it usable. Signed-off-by: mkrze <mateusz.krzemieniewski@intel.com>
CUDA allocates every copied output with at::empty(tensor.sizes(), ...), so copies are always contiguous. Tensor.to preserves the input layout, so transposed and channels-last inputs came back non-contiguous on XPU. Request a contiguous result. Tensors already on the target are still returned unchanged, as on CUDA. Add a two-device test with transposed, channels-last and row-strided inputs. Signed-off-by: mkrze <mateusz.krzemieniewski@intel.com>
The test deleted its inputs after the call but never reused the freed memory, so it could not detect a lifetime bug. Making it reuse the memory does not help on current hardware either: kernels from two streams on one XPU run in submission order on PVC and BMG, so a variant that copies on the wrong stream still produced correct results. Keep the test as a cross-device copy on non-default streams and drop the lifetime claim from its name and docstring. Signed-off-by: mkrze <mateusz.krzemieniewski@intel.com>
FBGEMM's merge_pooled_embeddings_test.py has test_all_to_one_device, which the plugin-local spread-over-devices test had copied. Enable that upstream test on the current accelerator in the shared patch, following reorder_batched_test.py: the class-level CUDA skip moves to the individual tests, so the merge_pooled_embeddings and sum_reduce_to_one tests stay CUDA-only. Add the module to test-fbgemm. Drop the copied test. The plugin-local tests now cover only what the upstream test does not check: views already on the target, contiguous copies, invalid devices and non-default streams. Signed-off-by: mkrze <mateusz.krzemieniewski@intel.com>
dvrogozh
pushed a commit
that referenced
this pull request
Oct 7, 2026
## Problem Several fbgemm-xpu sparse operators submitted their kernels to the SYCL queue of the *current* device instead of the device that holds their tensors: - `block_bucketize_sparse_features` / `block_bucketize_sparse_features_inference` - `populate_bucketized_permute` - `permute_1D_sparse_data`, `permute_2D_sparse_data` - `invert_permute` - `jagged_index_select_2d_forward` When the inputs live on another device, the kernel runs unordered with the pending work on the data's device and can read inputs before they are written. TorchRec row-wise QUANT inference hits this whenever its output device is not the current device: on an eight-card Data Center GPU Max host the bucketized indices were wrong in about half of the forwards. `expand_into_jagged_permute` was not affected: it already asked for the stream of its input's device explicitly. The same wrappers also checked only that inputs were XPU tensors, not that they were on the *same* XPU, so mixed-device inputs could pass pointers from another device to the selected queue and fail asynchronously. ## Changes - **Device guards.** Each operator now uses `SYCL_DEVICE_GUARD` on the tensor that FBGEMM's CUDA implementation guards, at the same point: `lengths` in the shared bucketize helper and in `populate_bucketized_permute`, `indices` in permute 1D/2D, `permute` in invert/expand permute, `values` in `jagged_index_select_2d_forward`. This is the pattern the jagged and bounds-check kernels already use. `expand_into_jagged_permute` now relies on the guard and the current stream too, instead of selecting the device explicitly. - **Same-device checks.** Like the CUDA wrappers (`TENSORS_ON_SAME_CUDA_GPU_IF_NOT_OPTIONAL`), bucketize, `populate_bucketized_permute`, permute 1D/2D (including optional `weights`) and `jagged_index_select_2d_forward` now call `TENSORS_ON_SAME_SYCL_XPU_IF_NOT_OPTIONAL` before selecting the queue. `expand_into_jagged_permute` already had an equivalent check; `invert_permute` has a single tensor. - **Behaviour change:** mixed-device inputs now raise `RuntimeError: Not all tensors were on the same XPU` instead of being accepted. - **Tests:** new `tests/test_non_current_device.py` with 15 two-device cases: - 8 correctness cases: each operator runs on `xpu:1` while `xpu:0` is current, with inputs that are still waiting for slow work on `xpu:1`, compared exactly with FBGEMM's CPU kernels; - 7 rejection cases for mixed-device inputs, including optional `weights` on the wrong device. - **Test fixture:** `test_int_nbit_lookup.py` accepted only torch 2.14.0, but `torch~=2.14.0` now installs 2.14.1, which turned every test in that file into a setup error in CI. The fixture now accepts any 2.14.x (same change as in #140). ## CI coverage The bug only shows up with two XPU devices (data on a device other than the current one), so it cannot be reproduced on one GPU. Upstream FBGEMM does not test this case either. **PR CI has a single GPU, so it collects the 15 new tests and skips them** ("requires two XPU devices"), the same way it already skips the two-device cases in `test_jagged_device_validation.py`. A green CI run therefore does not exercise this fix. The tests are meant for multi-GPU validation, for example a nightly run on a PVC machine with two cards. They need no extra setup and run as part of the normal test suite. To run them by hand on such a machine: ZE_FLAT_DEVICE_HIERARCHY=FLAT python -m pytest -rs packages/fbgemm-xpu/tests/test_non_current_device.py ## Testing - **PVC (Data Center GPU Max), two devices, PyTorch 2.14.1, manual:** 15/15 passed (repeated runs). On `main` the correctness case fails for every affected operator and passes for `expand_into_jagged_permute`, which was already correct. - **BMG (Arc Pro B60), PyTorch 2.14.1, same setup as CI:** plugin suite 340 passed, 28 skipped (15 of the skips are the new two-device cases). The patched upstream FBGEMM modules for these operators pass (`expand_into_jagged_permute`, `jagged_index_select_2d`, `block_bucketize`, `misc_ops`, `permute_sparse_features`). - **Performance (BMG):** per-call latency of the affected operators, `main` vs this branch, median of 3 alternating rounds: no change beyond run-to-run noise. Re-measured for `expand_into_jagged_permute` after its change. - `ruff check` and `bandit -r` are clean. --------- Signed-off-by: mkrze <mateusz.krzemieniewski@intel.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
3 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.
TorchRec's
SeqEmbeddingsAllToOneusesfbgemm::all_to_one_deviceto gather sequence embedding outputs from several devices onto one. Without an XPU implementation, shardedQuantEmbeddingCollectioninference across XPU devices fails withNotImplementedError.This PR adds a host-side XPU implementation:
The pooled merge operators (
merge_pooled_embeddingsandsum_reduce_to_one) remain unimplemented on XPU.Testing uses a hybrid approach. FBGEMM's upstream
all_to_one_devicetest now runs on XPU through the shared patch and checks values for inputs distributed across devices. Plugin-local tests cover the additional XPU contracts that upstream does not check: same-device aliasing, invalid inputs, contiguous copied outputs and copies on non-default streams. Multi-device tests require two XPUs and are skipped on single-GPU CI.This PR also updates the existing
test_int_nbit_lookup.pyfixture to accept any PyTorch 2.14.x release, matching thetorch~=2.14.0dependency. The previous exact 2.14.0 check caused all tests in that file to fail during setup after CI started installing 2.14.1.Validation: