build.toml

[general]
name = "punica_sgmv"

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

[kernel.sgmv]
language = "cuda"
src = [
  "sgmv/sgmv_cutlass.cu",
  "sgmv/sgmv_cutlass.cuh",
]
depends = [ "cutlass_3_8", "torch" ]

[kernel.sgmv_flashinfer]
language = "cuda"
cuda-capabilities = [ "8.0", "8.6", "8.7", "8.9", "9.0", "10.0", "10.1", "12.0" ]
src = [
  "flashinfer/cp_async.cuh",
  "flashinfer/mma.cuh",
  "flashinfer/permuted_smem.cuh",
  "flashinfer/vec_dtypes.cuh",
  "sgmv_flashinfer/sgmv_all.cu",
  "sgmv_flashinfer/sgmv_config.h",
  "sgmv_flashinfer/sgmv_flashinfer.cuh"
]
include = [ "." ]
depends = [ "torch" ]

[kernel.bgmv]
language = "cuda"
cuda-capabilities = [ "8.0", "8.6", "8.7", "8.9", "9.0", "10.0", "10.1", "12.0" ]
src = [
  "bgmv/bgmv_all.cu",
  "bgmv/bgmv_impl.cuh",
  "bgmv/bgmv_config.h",
  "flashinfer/vec_dtypes.cuh"
]
include = [ "." ]
depends = [ "torch" ]

[kernel.punica_kernels]
language = "cuda"
src = [
  "bgmv/bgmv_config.h",
  "punica_kernels/punica_ops.cc",
  "sgmv/sgmv.h",
  "sgmv_flashinfer/sgmv_config.h"
]
include = [ "." ]
depends = [ "torch" ]