Skip to content

[CUDA] Expose compilation targets via device_info - #4761

Open
dhiltgen wants to merge 1 commit into
ml-explore:mainfrom
dhiltgen:codex/cuda-architecture-info
Open

dhiltgen wants to merge 1 commit into
ml-explore:mainfrom
dhiltgen:codex/cuda-architecture-info

Conversation

@dhiltgen

Copy link
Copy Markdown
Contributor

This change exposes the set of CUDA compilation targets via device info so clients can determine if the build supports the discovered GPU's compute capability. This does not filter out unsupported devices to avoid ID reordering challenges if multiple devices are present, where only some are supported.

Clients can implement the compatibility logic outlined in https://docs.nvidia.com/cuda/cuda-programming-guide/01-introduction/cuda-platform.html#binary-compatibility and https://docs.nvidia.com/cuda/cuda-programming-guide/01-introduction/cuda-platform.html#ptx-compatibility to skip over unsupported devices.

Hypothetical example on a Blackwell GPU when compiled with MLX_CUDA_ARCHITECTURES=75-real;80-real;120a-real;120-virtual

{
    ...
    "architecture": "sm_120",
    "cuda_compilation_targets": "sm_75;sm_80;sm_120a;compute_120",
    ...
}
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: co-developed with coding agent

Comment thread mlx/backend/cuda/CMakeLists.txt Outdated
endif()
endforeach()
list(JOIN MLX_CUDA_COMPILATION_TARGETS ";" MLX_CUDA_COMPILATION_TARGETS)
configure_file(cuda_build_info.h.in cuda_build_info.h @ONLY)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

configure_file complicates build a lot, for simple case like this I prefer using target_compile_definitions to add a define directly.

Comment thread mlx/backend/cuda/device_info.cpp Outdated
// Compute capability as size_t (to match Metal's variant type)
dev.info["compute_capability_major"] = static_cast<size_t>(prop.major);
dev.info["compute_capability_minor"] = static_cast<size_t>(prop.minor);
dev.info["cuda_compilation_targets"] = cu::compilation_targets;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is no really a device info, it is independent of device. I think we can mirror the torch.cuda.get_arch_list API with mx.cuda.get_arch_list.

@dhiltgen
dhiltgen force-pushed the codex/cuda-architecture-info branch from 61481cf to 496c115 Compare October 26, 2026 05:02
@dhiltgen

Copy link
Copy Markdown
Contributor Author

Updated to use an explicit API instead of putting into device_info.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants