EnzymeAD / EnzymeAD/Enzyme-JAX
Attention Ops
Open
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
## Op Specification
- https://iree.dev/reference/mlir-dialects/LinalgExt/#iree_linalg_extattention-linalgextattentionop
- https://iree.dev/reference/mlir-dialects/LinalgExt/#iree_linalg_extonline_attention-linalgextonlineattentionop
## Lowering
- Enables us to lower directly to cudnn like https://github.com/jax-ml/jax/blob/da106b971af143906c8cbec8f4f2dfb7b57a11ce/jax/_src/cudnn/fused_attention_stablehlo.py#L252
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.