Repository: pytorch
Issue: Test test_vjp_nn_functional_scaled_dot_product_attention_cuda_float32 fails with `query: last dimension must be contiguous` on H100
Labels: triaged,module: functorch,module: sdpa

### 🐛 Describe the bug

Test `TestOperatorsCUDA.test_vjp_nn_functional_scaled_dot_product_attention_cuda_float32` fails with `query: last dimension must be contiguous` on **H100**

Also the schema test is still failing https://github.com/pytorch/pytorch/issues/102029#issuecomment-1558278905

```python
root@afe28e73ccd4:/opt/pytorch/pytorch# TORCH_SHOW_CPP_STACKTRACES=1 python test/functorch/test_ops.py -v -k TestOperatorsCUDA.test_vjp_nn_functional_scaled_dot_product_attention_cuda_float32
test_vjp_nn_functional_scaled_dot_product_attention_cuda_float32 (__main__.TestOperatorsCUDA) ... /opt/pytorch/pytorch/torch/_functorch/deprecated.py:73: UserWarning: We've integrated functorch into PyTorch. As the final step of the integration, functorch.vjp is deprecated as of PyTorch 2.0 and will be deleted in a future version of PyTorch >= 2.3. Please use torch.func.vjp instead; see the PyTorch 2.0 release notes and/or the torch.func migration guide for more details https://pytorch.org/docs/master/func.migrating.html
  warn_deprecated('vjp')
ERROR

======================================================================
ERROR: test_vjp_nn_functional_scaled_dot_product_attention_cuda_float32 (__main__.TestOperatorsCUDA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/opt/pytorch/pytorch/torch/testing/_internal/common_utils.py", line 2179, in wrapper
    method(*args, **kwargs)
  File "/opt/pytorch/pytorch/torch/testing/_internal/common_device_type.py", line 427, in instantiated_test
    raise rte
  File "/opt/pytorch/pytorch/torch/testing/_internal/common_device_type.py", line 414, in instantiated_test
    result = test(self, **param_kwargs)
  File "/opt/pytorch/pytorch/torch/testing/_internal/common_utils.py", line 1154, in wrapper
    fn(*args, **kwargs)
  File "/opt/pytorch/pytorch/torch/testing/_internal/common_device_type.py", line 904, in test_wrapper
    return test(*args, **kwargs)
  File "/opt/pytorch/pytorch/test/functorch/test_ops.py", line 670, in test_vjp
    _test(op)
  File "/opt/pytorch/pytorch/test/functorch/test_ops.py", line 660, in _test
    out_noncontig, vjp_fn = vjp(noncontig_fn, *noncontig_primals)
  File "/opt/pytorch/pytorch/torch/_functorch/deprecated.py", line 74, in vjp
    return _impl.vjp(func, *primals, has_aux=has_aux)
  File "/opt/pytorch/pytorch/torch/_functorch/eager_transforms.py", line 264, in vjp
    return _vjp_with_argnums(func, *primals, has_aux=has_aux)
  File "/opt/pytorch/pytorch/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "/opt/pytorch/pytorch/torch/_functorch/eager_transforms.py", line 291, in _vjp_with_argnums
    primals_out = func(*primals)
  File "/opt/pytorch/pytorch/test/functorch/test_ops.py", line 128, in wrapped
    result = f(*_args, **kwargs)
  File "/opt/pytorch/pytorch/torch/testing/_internal/opinfo/core.py", line 1097, in __call__
    return self.op(*args, **kwargs)
  File "/opt/pytorch/pytorch/torch/testing/_internal/common_methods_invocations.py", line 13237, in <lambda>
    wrapper_set_seed(torch.nn.functional.scaled_dot_product_attention, *args, **kwargs),
  File "/opt/pytorch/pytorch/torch/testing/_internal/common_methods_invocations.py", line 8906, in wrapper_set_seed
    return op(*args, **kwargs)
RuntimeError: query: last dimension must be contiguous
Exception raised from _efficient_attention_forward at /opt/pytorch/pytorch/aten/src/ATen/native/transformers/cuda/attention.cu:862 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0xae (0x7f91acbb014e in /opt/pytorch/pytorch/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, char const*) + 0x7d (0x7f91acb65cdb in /opt/pytorch/pytorch/torch/lib/libc10.so)
frame #2: at::native::_efficient_attention_forward(at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, c10::optional<long>, double, long, bool, c10::optional<double>, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&) + 0x2b55 (0x7f9159594d35 in /opt/pytorch/pytorch/torch/lib/libtorch_cuda.so)
frame #3: <unknown function> + 0x2d24a70 (0x7f91597aaa70 in /opt/pytorch/pytorch/torch/lib/libtorch_cuda.so)
frame #4: <unknown function> + 0x2d24ba7 (0x7f91597aaba7 in /opt/pytorch/pytorch/torch/lib/libtorch_cuda.so)
frame #5: at::_ops::_efficient_attention_forward::call(at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, c10::optional<long>, double, long, bool, c10::optional<double>, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&) + 0x295 (0x7f917834be75 in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #6: at::native::_scaled_dot_product_efficient_attention_cuda(at::Tensor const&, at::Tensor const&, at::Tensor const&, bool, bool, c10::optional<double>) + 0x19b (0x7f915955ee3b in /opt/pytorch/pytorch/torch/lib/libtorch_cuda.so)
frame #7: <unknown function> + 0x2d2414a (0x7f91597aa14a in /opt/pytorch/pytorch/torch/lib/libtorch_cuda.so)
frame #8: <unknown function> + 0x2d241e1 (0x7f91597aa1e1 in /opt/pytorch/pytorch/torch/lib/libtorch_cuda.so)
frame #9: at::_ops::_scaled_dot_product_efficient_attention::redispatch(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor const&, bool, bool, c10::optional<double>) + 0xd1 (0x7f917828c401 in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #10: <unknown function> + 0x3af504f (0x7f9179ca204f in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #11: <unknown function> + 0x3af5cc9 (0x7f9179ca2cc9 in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #12: <unknown function> + 0x11c2bb7 (0x7f917736fbb7 in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #13: <unknown function> + 0x11c31e9 (0x7f91773701e9 in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #14: at::functorch::GradInterpreterPtr::sendToNextInterpreterImpl(c10::OperatorHandle const&, std::vector<c10::IValue, std::allocator<c10::IValue> >*, bool) + 0x3b (0x7f917737093b in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #15: <unknown function> + 0x12cc42a (0x7f917747942a in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #16: <unknown function> + 0x12c7a8b (0x7f9177474a8b in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #17: at::_ops::_scaled_dot_product_efficient_attention::redispatch(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor const&, bool, bool, c10::optional<double>) + 0x21f (0x7f917828c54f in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #18: <unknown function> + 0x3af504f (0x7f9179ca204f in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #19: <unknown function> + 0x3af5bc4 (0x7f9179ca2bc4 in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #20: at::_ops::_scaled_dot_product_efficient_attention::call(at::Tensor const&, at::Tensor const&, at::Tensor const&, bool, bool, c10::optional<double>) + 0x1ae (0x7f917834b4be in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #21: at::native::scaled_dot_product_attention(at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, double, bool, c10::optional<double>) + 0x116a (0x7f9177d3951a in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #22: <unknown function> + 0x2709fc9 (0x7f91788b6fc9 in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #23: <unknown function> + 0x11c2bb7 (0x7f917736fbb7 in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #24: <unknown function> + 0x11c3a5d (0x7f9177370a5d in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #25: <unknown function> + 0x12cc2a7 (0x7f91774792a7 in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #26: <unknown function> + 0x12c58ac (0x7f91774728ac in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #27: at::_ops::scaled_dot_product_attention::call(at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, double, bool, c10::optional<double>) + 0x487 (0x7f91781d4427 in /opt/pytorch/pytorch/torch/lib/libtorch_cpu.so)
frame #28: <unknown function> + 0x5da5a6 (0x7f91810fe5a6 in /opt/pytorch/pytorch/torch/lib/libtorch_python.so)
<omitting python frames>


----------------------------------------------------------------------
Ran 1 test in 0.631s

FAILED (errors=1)
```

### Versions

cuda 12.1
H100
torch_commit: https://github.com/pytorch/pytorch/commit/7c2641d5f1081811e664267406df1059687aad7a

test failures started on 5/19/23
see also #102029

cc @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg @mikaylagawarecki @zou3519 @Chillee @samdow @kshitij12345 @janeyx99 @Fuzzkatt @ptrblck @ngimel