Skip to content

Commit

Permalink
Lift jagged -> padded dense forward / backward kernels from fbgemm_gpu (
Browse files Browse the repository at this point in the history
pytorch#125946)

PyTorch can't depend on `fbgemm_gpu` as a dependency because `fbgemm_gpu` already has a dependency on PyTorch. So this PR copy / pastes kernels from `fbgemm_gpu`:
* `dense_to_jagged_forward()` as CUDA registration for new ATen op `_padded_dense_to_jagged_forward()`
* `jagged_to_padded_dense_forward()` as CUDA registration for new ATen op `_jagged_to_padded_dense_forward()`

CPU impls for these new ATen ops will be added in a follow-up PR.
Pull Request resolved: pytorch#125946
Approved by: https://github.com/davidberard98
  • Loading branch information
jbschlosser authored and petrex committed Jun 5, 2024
1 parent 61690a5 commit ef3a552
Show file tree
Hide file tree
Showing 4 changed files with 1,118 additions and 0 deletions.
10 changes: 10 additions & 0 deletions aten/src/ATen/native/native_functions.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -14644,6 +14644,16 @@
NestedTensorCUDA: NestedTensor_to_padded_tensor_cuda
autogen: to_padded_tensor.out

- func: _jagged_to_padded_dense_forward(Tensor values, Tensor[] offsets, SymInt[] max_lengths, float padding_value=0.0) -> Tensor
variants: function
dispatch:
CUDA: _fbgemm_jagged_to_padded_dense_forward

- func: _padded_dense_to_jagged_forward(Tensor dense, Tensor[] offsets, SymInt? total_L=None) -> Tensor
variants: function
dispatch:
CUDA: _fbgemm_dense_to_jagged_forward_symint

- func: _nested_tensor_softmax_with_shape(Tensor self, Tensor query) -> Tensor
dispatch:
NestedTensorCPU: NestedTensor_softmax_dropout
Expand Down
Loading

0 comments on commit ef3a552

Please sign in to comment.