build.toml

[general]
name = "residual_rms_rocm"
universal = false

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

[kernel.residual_rms_rocm]
depends = ["torch"]
backend = "rocm"
rocm-archs = [
    "gfx90a",
]
src = [
    "residual_rms_rocm/residual_rms_dispatch.cu",
    "residual_rms_rocm/residual_rms_scalar.cu",
    "residual_rms_rocm/residual_rms_vectorized.cu",
    "residual_rms_rocm/utils.h",
]
include = ["."]