build.toml

[general]
name = "adam_atan2"
universal = false

[torch]
src = [
  "torch-ext/torch_binding.cpp",
  "torch-ext/torch_binding.h"
]

[kernel.adam_atan2]
backend = "cuda"
cuda-capabilities = [
    "7.0",
    "7.2",
    "7.5",
    "8.0",
    "8.6",
    "8.7",
    "8.9",
    "9.0",
    "10.0",
    "10.1",
    "11.8",
    "12.0"
]
depends = ["torch", "cutlass_3_8"]
src = [
    "csrc/adam_atan2.cu",
]