Skip to content

nix-builder: use ptxas 13.4 for CUDA 13.0 builds on aarch64-linux #952

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
drbh wants to merge 1 commit into from
Closed

nix-builder: use ptxas 13.4 for CUDA 13.0 builds on aarch64-linux #952

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 37 additions & 0 deletions nix-builder/lib/checks.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,43 @@ let
}
];
};
testFlattenSystemsPtxasVersion = {
expr = flattenSystems [ (validVersion // { ptxasVersion = "13.4"; }) ];
expected = [
{
torchVersion = "2.13";
cudaVersion = "13.0";
ptxasVersion = "13.4";
system = "x86_64-linux";
bundleBuild = true;
}
{
torchVersion = "2.13";
cudaVersion = "13.0";
ptxasVersion = "13.4";
system = "aarch64-linux";
bundleBuild = true;
}
];
};
testFlattenSystemsPtxasVersionPerSystem = {
expr = flattenSystems [ (validVersion // { ptxasVersion.aarch64-linux = "13.4"; }) ];
expected = [
{
torchVersion = "2.13";
cudaVersion = "13.0";
system = "x86_64-linux";
bundleBuild = true;
}
{
torchVersion = "2.13";
cudaVersion = "13.0";
ptxasVersion = "13.4";
system = "aarch64-linux";
bundleBuild = true;
}
];
};
testFlattenSystemsValidates = {
expr = throws (flattenSystems [ (validVersion // { cudaVerison = "13.0"; }) ]);
expected = true;
Expand Down
7 changes: 5 additions & 2 deletions nix-builder/lib/mk-build-set.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,14 @@ let
self: super: {
cuda_nvcc = super.cuda_nvcc.overrideAttrs (prevAttrs: {
# Do this before the original postInstall, so that subsequent
# postInstall steps are applied.
# postInstall steps are applied. Since CUDA 13, nvvm is not
# part of cuda_nvcc anymore, but a separate libnvvm package.
postInstall = ''
cp ${ptxasPackages.cuda_nvcc.src}/bin/ptxas $out/bin/ptxas
cp ${ptxasPackages.cuda_nvcc.src}/bin/nvlink $out/bin/nvlink
cp ${ptxasPackages.cuda_nvcc.src}/nvvm/bin/* $out/nvvm/bin/
if [ -d ${ptxasPackages.cuda_nvcc.src}/nvvm/bin ]; then
cp ${ptxasPackages.cuda_nvcc.src}/nvvm/bin/* $out/nvvm/bin/
fi
''
+ prevAttrs.postInstall or "";
});
Expand Down
26 changes: 25 additions & 1 deletion nix-builder/lib/torch-version-utils.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,13 +41,37 @@ rec {
version
);

# `ptxasVersion` is either a version string that applies to all systems,
# or an attrset { <system> = <version>; } to override ptxas only on some
# systems. Resolve it for a single system; `null` means: use the ptxas
# from `cudaVersion`.
ptxasVersionForSystem =
version: system:
let
ptxasVersion = version.ptxasVersion or null;
in
if builtins.isAttrs ptxasVersion then ptxasVersion.${system} or null else ptxasVersion;

# Expand { systems = [ a b ]; .. } to [ { system = a; ..} { system = b; .. } ]
flattenSystems =
versions:
lib.foldl' (
acc: version:
acc
++ map (system: (builtins.removeAttrs version [ "systems" ]) // { inherit system; }) version.systems
++ map (
system:
let
ptxasVersion = ptxasVersionForSystem version system;
in
(builtins.removeAttrs version [
"systems"
"ptxasVersion"
])
// {
inherit system;
}
// lib.optionalAttrs (ptxasVersion != null) { inherit ptxasVersion; }
) version.systems
) [ ] (map validateTorchVersion versions);

backend =
Expand Down
10 changes: 10 additions & 0 deletions nix-builder/versions.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@
{
torchVersion = "2.10";
cudaVersion = "13.0";

@danieldk danieldk Oct 29, 2026

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 should not be done. This switches all 13.0 builds to ptxas 13.4. It is a hack and should only be done on kernels that really need it. See

https://github.com/huggingface/kernels-community/blob/4498ed972091eb7d087fdb3853417f706193df28/flash-attn3/flake.nix#L26

# ptxas 13.0 segfaults on aarch64 compiling large sm90 kernels (flash-attn3).
ptxasVersion.aarch64-linux = "13.4";
systems = [
"x86_64-linux"
"aarch64-linux"
Expand Down Expand Up @@ -83,6 +85,8 @@
{
torchVersion = "2.11";
cudaVersion = "13.0";
# ptxas 13.0 segfaults on aarch64 compiling large sm90 kernels (flash-attn3).
ptxasVersion.aarch64-linux = "13.4";
systems = [
"x86_64-linux"
"aarch64-linux"
Expand Down Expand Up @@ -152,6 +156,8 @@
{
torchVersion = "2.12";
cudaVersion = "13.0";
# ptxas 13.0 segfaults on aarch64 compiling large sm90 kernels (flash-attn3).
ptxasVersion.aarch64-linux = "13.4";
systems = [
"x86_64-linux"
"aarch64-linux"
Expand Down Expand Up @@ -216,6 +222,8 @@
{
torchVersion = "2.13";
cudaVersion = "13.0";
# ptxas 13.0 segfaults on aarch64 compiling large sm90 kernels (flash-attn3).
ptxasVersion.aarch64-linux = "13.4";
systems = [
"x86_64-linux"
"aarch64-linux"
Expand Down Expand Up @@ -283,6 +291,8 @@
{
torchVersion = "2.14";
cudaVersion = "13.0";
# ptxas 13.0 segfaults on aarch64 compiling large sm90 kernels (flash-attn3).
ptxasVersion.aarch64-linux = "13.4";
systems = [
"x86_64-linux"
"aarch64-linux"
Expand Down
Loading