File size: 658 Bytes
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 |
[general]
name = "megablocks"
universal = false
[torch]
src = [
"torch-ext/torch_binding.cpp",
"torch-ext/torch_binding.h"
]
[kernel.megablocks]
backend = "cuda"
src = [
"csrc/new_cumsum.h",
"csrc/new_cumsum.cu",
"csrc/new_histogram.h",
"csrc/new_histogram.cu",
"csrc/new_indices.h",
"csrc/new_indices.cu",
"csrc/new_replicate.cu",
"csrc/new_replicate.h",
"csrc/new_sort.h",
"csrc/new_sort.cu",
]
depends = [ "torch", "cutlass_3_8" ]
[test]
python-git-packages = [
{ url = "https://github.com/stanford-futuredata/stk.git", rev = "7363137", sha256 = "0m6g5l9nlwaiwybg5j8dhnz159wdpabdnkzapnn3dsifxrsb59vz" }
] |