build.toml

[general]
name = "rwkv"
universal = false

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

[kernel.rwkv]
depends = ["torch"]
backend = "cuda"
cuda-capabilities = [
    "8.0",
    "8.9",
    "9.0",
    "10.0",
    "12.0",
]
include = ["."]
src = [
    "rwkv/wkv_cuda.cu",
    "rwkv/wkv_cuda_bf16.cu",
]
cuda-flags = [
        "-res-usage",
        "--use_fast_math",
        "-O3",
        "--extra-device-vectorization",
        "-DTmax=1024",
]