🚀 The feature, motivation and pitch

torch.utils.cpp_extension.CUDAExtension already supports relocatable device code / device linking workflows (for example, via dlink=True, dlink_libraries, and nvcc_dlink), but torch.utils.cpp_extension.load() does not expose an equivalent capability.

This creates a feature gap between the AOT/setuptools path and the JIT path.

Motivation:

• Some CUDA extensions require an extra device-link step (-dlink) before the final host link step.
• Typical examples include:
  • CUDA dynamic parallelism
  • multi-CU device symbol references
  • linking against static CUDA libraries that contain RDC-compiled objects
• These use cases are already recognized by PyTorch and documented for CUDAExtension, but they are not available through load() / JIT extensions.

Today, users who want the convenience of torch.utils.cpp_extension.load() cannot express the same device-linking workflow without patching cpp_extension.py locally or moving to a different build path.

I would like PyTorch to add support for CUDA device linking in torch.utils.cpp_extension.load() (and potentially load_inline() if appropriate), ideally with an API that is consistent with or analogous to CUDAExtension.

Related context:

• #57543 requested CUDA dynamic parallelism / device linking support for custom PyTorch CUDA extensions
• #78225 added relocatable device code linking support for CUDAExtension
• #44279 shows user confusion around how to achieve this in cpp extensions, especially for JIT-like workflows

Expected outcome:

• JIT cpp extensions should be able to trigger the extra CUDA device-link step when needed
• The behavior should remain opt-in and preserve current defaults for existing users

Alternatives

Current alternatives are not ideal:

1. Switch from torch.utils.cpp_extension.load() to the setuptools / CUDAExtension path.

   • This loses the convenience of JIT-style workflows.
   • It adds packaging/build-system overhead for users who only want runtime compilation.

2. Patch torch/utils/cpp_extension.py locally.

   • This is fragile and not maintainable across PyTorch upgrades.

3. Manually reimplement the build flow with custom ninja / cmake logic.

   • This defeats the purpose of using PyTorch's extension helpers.

Because CUDAExtension already supports this class of workflows, exposing similar functionality in load() seems like the most natural solution.

Additional context

This request is not asking for RDC/device-link support in general from scratch, because that already exists for CUDAExtension. Instead, this issue is specifically about feature parity for the JIT path:

• torch.utils.cpp_extension.load()
• potentially torch.utils.cpp_extension.load_inline() A possible implementation might expose:
• a higher-level option such as dlink=True
• or a lower-level way to pass nvcc device-link flags I am intentionally not prescribing the exact API here; the main request is to support the extra CUDA device-link step in the JIT extension flow.

cc @janeyx99 @ptrblck @msaroufim @eqy @jerryzh168 @tinglvv @nWEIdia