Skip to content

Commit e2f0818

Browse files
jbschlosserpytorchmergebot
authored andcommitted
Lift jagged -> padded dense forward / backward kernels from fbgemm_gpu (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
1 parent 3f5b59e commit e2f0818

File tree

4 files changed

+1108
-0
lines changed

4 files changed

+1108
-0
lines changed

aten/src/ATen/native/native_functions.yaml

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14636,6 +14636,16 @@
1463614636
NestedTensorCUDA: NestedTensor_to_padded_tensor_cuda
1463714637
autogen: to_padded_tensor.out
1463814638

14639+
- func: _jagged_to_padded_dense_forward(Tensor values, Tensor[] offsets, SymInt[] max_lengths, float padding_value=0.0) -> Tensor
14640+
variants: function
14641+
dispatch:
14642+
CUDA: _fbgemm_jagged_to_padded_dense_forward
14643+
14644+
- func: _padded_dense_to_jagged_forward(Tensor dense, Tensor[] offsets, SymInt? total_L=None) -> Tensor
14645+
variants: function
14646+
dispatch:
14647+
CUDA: _fbgemm_dense_to_jagged_forward_symint
14648+
1463914649
- func: _nested_tensor_softmax_with_shape(Tensor self, Tensor query) -> Tensor
1464014650
dispatch:
1464114651
NestedTensorCPU: NestedTensor_softmax_dropout

0 commit comments

Comments
 (0)