Interface with Flash Attention 2.0

ai4co / rl4co

A PyTorch library for all things Reinforcement Learning (RL) for Combinatorial Optimization (CO)

https://rl4.co

MIT License

455 stars 84 forks source link

Interface with Flash Attention 2.0 #85

Closed fedebotu closed 1 year ago

fedebotu commented 1 year ago

The newest Flash Attention version is out, it would be great to wrap it as optional with our API. Looks like there are several optimizations on it! Preprint Code

fedebotu commented 1 year ago

Implemented in https://github.com/kaist-silab/rl4co/releases/tag/v0.2.3