build.toml

[general]
name = "metal_flash_sdpa"
universal = false

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

[kernel.sdpa_metal]
backend = "metal"
src = [
  "sdpa-metal/scaled_dot_product_attention.mm",
  "sdpa-metal/scaled_dot_product_attention.metal",
  "sdpa-metal/common.h",
]
depends = [ "torch" ]