Repository: pytorch
Issue: CUDA Extension Bug for Pytorch 2.2.0
Labels: high priority,module: cpp-extensions,module: cuda,oncall: releng,triaged,topic: binaries

### 🐛 Describe the bug

Seems like the released Pytorch 2.2.0 for CUDA Linux has bugs to run C++ extensions. However, Pytorch 2.2.0 for CPU is not affected by this. So looks like this is a "2.2.0 + CUDA-only" bug.

To reproduce:
```sh
$ python3.8 -m pip install torch==2.2.0 --index-url https://download.pytorch.org/whl/cu118

$ python3.8 -m pip install --verbose --upgrade git+https://github.com/microsoft/tutel@main

$ python3.8 -m tutel.examples.helloworld --batch_size=16
terminate called after throwing an instance of 'c10::Error'
  what():  !dispatch_key_.has_value() INTERNAL ASSERT FAILED at "../aten/src/ATen/core/library.cpp":82, please report a bug to PyTorch. (Error occurred while processing TORCH_LIBRARY block at ./tutel/custom/custom_kernel.cpp:891)
Exception raised from Library at ../aten/src/ATen/core/library.cpp:82 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7ff953aded87 in /usr/local/lib/python3.8/dist-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x64 (0x7ff953a8f75f in /usr/local/lib/python3.8/dist-packages/torch/lib/libc10.so)
frame #2: c10::detail::torchInternalAssertFail(char const*, char const*, unsigned int, char const*, std::string const&) + 0x3f (0x7ff953adc8bf in /usr/local/lib/python3.8/dist-packages/torch/lib/libc10.so)
frame #3: torch::Library::Library(torch::Library::Kind, std::string, std::optional<c10::DispatchKey>, char const*, unsigned int) + 0x96c (0x7ff98b9a71bc in /usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so)
frame #4: <unknown function> + 0x189c6 (0x7ff8e0dd29c6 in /usr/local/lib/python3.8/dist-packages/tutel_custom_kernel.cpython-38-x86_64-linux-gnu.so)
frame #5: <unknown function> + 0x108d3 (0x7ff9a58348d3 in /lib64/ld-linux-x86-64.so.2)
frame #6: <unknown function> + 0x1539f (0x7ff9a583939f in /lib64/ld-linux-x86-64.so.2)
frame #7: _dl_catch_exception + 0x6f (0x7ff9a559a16f in /lib/x86_64-linux-gnu/libc.so.6)
frame #8: <unknown function> + 0x1496a (0x7ff9a583896a in /lib64/ld-linux-x86-64.so.2)
frame #9: <unknown function> + 0xf96 (0x7ff9a5010f96 in /lib/x86_64-linux-gnu/libdl.so.2)
frame #10: _dl_catch_exception + 0x6f (0x7ff9a559a16f in /lib/x86_64-linux-gnu/libc.so.6)
frame #11: _dl_catch_error + 0x2f (0x7ff9a559a1ff in /lib/x86_64-linux-gnu/libc.so.6)
frame #12: <unknown function> + 0x1745 (0x7ff9a5011745 in /lib/x86_64-linux-gnu/libdl.so.2)
frame #13: dlopen + 0x71 (0x7ff9a5011051 in /lib/x86_64-linux-gnu/libdl.so.2)
<omitting python frames>
frame #16: python3() [0x664d58]
frame #17: python3() [0x5c08c9]
frame #32: python3() [0x5ee6f0]
frame #37: python3() [0x4c6e27]
frame #38: python3() [0x5c08c9]
frame #51: python3() [0x5ee6f0]
frame #56: python3() [0x4c6e27]
frame #57: python3() [0x5c08c9]
```


If I turn to use Pytorch 2.0.0 or 2.1.0 with `python3 -m pip install https://download.pytorch.org/whl/cu118/torch-2.1.0%2Bcu118-cp38-cp38-linux_x86_64.whl` which is an all-in-one package, everything is fine.

BTW, I also tried other C++ Extensions, Pytorch 2.2.0 for CUDA always produces the same error.

### Versions

Verified that it is not related to Python version. This bug happens in both Python3.8 and Python3.12.

cc @ezyang @gchanan @zou3519 @kadeng @malfet @ptrblck