first version of flash_attention for jax #19743
Conversation
|
Thanks for the PR! Have you tried to time it on GPU compared to regular attention? I was under the impression that we were going to need a custom Pallas kernel for this. |
|
I have used /keras/src/layers/attention/ directory as a template for implementing a flash attention but I don't understand how the mask is generated in the benchmark. I need one but I don't see it |
|
Hi @fchollet Can you please review this PR? Thank you! |
| @@ -76,6 +82,44 @@ def relu6(x): | |||
| return Relu6().symbolic_call(x) | |||
| return backend.nn.relu6(x) | |||
|
|
|||
| @keras_export(["keras.ops.flash_attention", "keras.ops.nn.flash_attention"]) | |||
There was a problem hiding this comment.
Hello there,
That's a very wonderful addition. I've a bit of doubt there:
I suppose ops.flash_attention is an operation, while nn.flash_attention is just a neural network layer.
The basic difference between these two is that - an operation may not have any trainable parameter with it, while a neural network layer should have trainable parameters.
Am I right till now?
If yes, please provide separate examples of each one of them in the docs!
Best Regards,
Abhas Kumar Sinha
| Returns: | ||
| A tensor with the same shape as `x`. | ||
|
|
||
| Example: |
There was a problem hiding this comment.
Enclose the Example with "```".
"""python
This is an example documentation.
Example:
```
example = example()
```
"""This helps automated doc renderers to automatically find out code examples from the program docs and render those parts accordingly.
This is my first version of the flash attention implementation .It is just for Jax.