[package]
name = "gllm-kernels"
version = "0.1.3"
edition = "2021"
authors = ["gllm contributors"]
license = "Apache-2.0"
description = "Low-level attention kernels for gllm with CUDA/ROCm support"
repository = "https://github.com/putao520/gllm-kernels"
homepage = "https://github.com/putao520/gllm-kernels"
keywords = ["attention", "kernels", "burn", "gpu", "tensor"]
categories = ["algorithms", "science"]

[lib]
name = "gllm_kernels"
path = "src/lib.rs"

[features]
default = ["cpu"]
cpu = ["burn-ndarray"]
cuda = ["burn-cuda"]
wgpu = ["burn-wgpu"]
fusion = ["burn/fusion", "burn-fusion"]
nccl = ["cudarc/nccl", "cuda"]
cuda-kernel = ["cudarc", "half"]
rocm-kernel = ["half"]
flash-attention-v3 = [
    "flash-attention-v3-wgmma",
    "flash-attention-v3-async",
    "flash-attention-v3-fp8",
    "flash-attention-v3-block-quant",
]
flash-attention-v3-wgmma = []
flash-attention-v3-async = []
flash-attention-v3-fp8 = []
flash-attention-v3-block-quant = []
# hip 等 burn 支持后再加

[dependencies]
burn = { version = "0.20.0-pre.6", default-features = false }
burn-ndarray = { version = "0.20.0-pre.6", optional = true }
burn-cuda = { version = "0.20.0-pre.6", optional = true }
burn-wgpu = { version = "0.20.0-pre.6", optional = true }
burn-fusion = { version = "0.20.0-pre.6", optional = true }
log = "0.4"
crossbeam-channel = "0.5"
cudarc = { version = "0.18", optional = true, default-features = false, features = ["std", "driver", "nvrtc", "f16", "cuda-12090", "dynamic-loading"] }
half = { version = "2.7", optional = true }

[dev-dependencies]
criterion = "0.5"

[[bench]]
name = "attention_bench"
harness = false
