Skip to content

Implemention of philox key fold in tensor - #5451

Draft
orrangetabby17 wants to merge 3 commits into
mainfrom
shuxin/sycl/philox-key-fold-in-tensor
Draft

orrangetabby17 wants to merge 3 commits into
mainfrom
shuxin/sycl/philox-key-fold-in-tensor

Conversation

@orrangetabby17

@orrangetabby17 orrangetabby17 commented Sep 20, 2026

Copy link
Copy Markdown
Contributor

Resolves #4062

Summary

  1. This PR adds XPU implementation for _philox_key_fold_in.Tensor which is the tensor overload of _philox_key_fold_in.

  2. The implementation also reused the Philox fold-in logic into a shared helper so that the scalar and tensor variants use the same key folding logic.

  3. Add input validation for shape, dtype, device, and number of elements.

@github-actions github-actions Bot added disable_e2e Disable all e2e test jobs for the PR disable_distributed Disable distributed UT test jobs for the PR labels Sep 20, 2026
@orrangetabby17 orrangetabby17 changed the title Shuxin/sycl/philox key fold in tensor Implemention of philox key fold in tensor Sep 20, 2026
Port the graph-safe tensor-data overload of _philox_key_fold_in to XPU,
following the CUDA implementation in aten/src/ATen/native/cuda/PhiloxKeySplit.cu.
The data value is read from device memory at kernel execution time instead of
being baked into the launch. The scalar kernel body is factored into a shared
inline helper so both variants fold keys identically.

(cherry picked from commit 9858ab5)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

disable_distributed Disable distributed UT test jobs for the PR disable_e2e Disable all e2e test jobs for the PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Stateless RNG APIs for uniform & normal generation on XPU

1 participant