File size: 1,211 Bytes
2595c46 359242d 2595c46 359242d 2595c46 359242d 2595c46 359242d 2595c46 |
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 51 52 53 |
{
description = "Flake for megablocks_moe kernel";
inputs = {
kernel-builder.url = "github:huggingface/kernel-builder/support-custom-python-libraries-in-dev-shell-nixland";
# Add libraries as inputs
composer = {
url = "github:mosaicml/composer";
flake = false;
};
stk = {
url = "github:stanford-futuredata/stk";
flake = false;
};
# TODO: update to build with the correct torch version
# grouped_gemm = {
# url = "github:tgale96/grouped_gemm";
# flake = false;
# };
};
outputs = {
self,
kernel-builder,
composer,
stk,
# grouped_gemm,
}:
kernel-builder.lib.genFlakeOutputs {
path = ./.;
rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;
# Map custom packages to their sources
customPythonPackages = {
composer = composer;
stk = stk;
# grouped_gemm = grouped_gemm;
};
pythonTestDeps = [
"tqdm"
"py-cpuinfo"
"importlib-metadata"
"torchmetrics"
"composer"
"stk"
# "grouped_gemm"
# "yahp" # may be needed for some testing plugin
];
};
}
|