File size: 1,261 Bytes
eb8ddce 2e75662 eb8ddce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
{
description = "Flake for Hopper Flash Attention kernel";
inputs = {
kernel-builder.url = "github:huggingface/kernel-builder";
};
outputs =
{
self,
kernel-builder,
}:
kernel-builder.lib.genFlakeOutputs {
path = ./.;
rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;
# Building with CDUA later than 12.4 fails with:
#
# error: 'ptxas' died due to signal 11 (Invalid memory reference)
#
# So, build for 12.4 only and copy to all the other build variants
# by hand (which works fine thanks to backward compat).
torchVersions = [
{
torchVersion = "2.6";
cudaVersion = "12.4";
cxx11Abi = false;
systems = [ "x86_64-linux" ];
upstreamVariant = true;
}
{
torchVersion = "2.6";
cudaVersion = "12.4";
cxx11Abi = true;
systems = [ "x86_64-linux" ];
upstreamVariant = true;
}
{
torchVersion = "2.7";
cudaVersion = "12.4";
cxx11Abi = true;
systems = [
"x86_64-linux"
"aarch64-linux"
];
upstreamVariant = true;
}
];
};
}
|