❓ [Question] Exporting models using FlashAttention package

[s1ddok]

I'd love to export a PyTorch model to TensorRT. In this model I use flash-attn package to speed-up attention. Is this supported?

[narendasan]

Im not sure how much of a difference it will make vs the kernels TensorRT will pick automatically but if you want to use a specific kernel you can through plugins: https://docs.pytorch.org/TensorRT/tutorials/_rendered_examples/dynamo/aot_plugin.html?highlight=aot+plugin. This system can be used with any PTX source. Alternatively, you can also make a CUDA/C++ based plugin which will be maximally performant.