Explore potential speed benefits from implementing kernl (Up to 12X faster GPU inference)

[0xdevalias]

Is your feature request related to a problem? Please describe.
N/A

Describe the solution you'd like
Explore whether implementing kernl would provide speedups for Stable Diffusion.

• https://github.com/ELS-RD/kernl

  • Kernl lets you run PyTorch transformer models several times faster on GPU with a single line of code, and is designed to be easily hackable.

• https://www.reddit.com/r/MachineLearning/comments/ydqmjp/p_up_to_12x_faster_gpu_inference_on_bert_t5_and/

  • Up to 12X faster GPU inference on Bert, T5 and other transformers with OpenAI Triton kernels

Describe alternatives you've considered

• That kernl won't be useful in this repo, or that existing optimisations already do things as well as it could (or better).

Additional context

• Crossposted on:
  • [Feature Request]: Explore potential StableDiffusion speed benefits from implementing kernl (Up to 12X faster GPU inference) AUTOMATIC1111/stable-diffusion-webui#4096
  • Would it be possible to use kernl to speed up Stable Diffusion? ELS-RD/kernl#141

[patrickvonplaten]

Very interesting! cc @NouamaneTazi @Narsil @patil-suraj here

[patrickvonplaten]

I'm not sure the diffusers team has time for this, but we'd definitely be interested in exploring this option, but we'd need some benchmark numbers :-)

[Narsil]

@OlivierDehaene Played with it and saw same or even slightly worse performance on some hardware.

Not to diminish the acomplishments of kernl, I am a big fan of putting kernel writing in the hands of users, and being able to tinker with it.
I am just saying running some benchmarks would be necessary. It would be awesome if we manage to expose a simple script.

[OlivierDehaene]

I played a bit with their Linear, LayerNorm and FlashAttention implementations with our optimized BLOOM deployement:

• I didn't see gains with their Linaer and LayerNorm but maybe I did something wrong I didn't have the time to investigate more.
• Flash Attention is awesome as always and their implementation adds support for arbitrary masks not just causal masks which is great :) It is also 100% Triton which makes it way easier to read/import/modify.

[patil-suraj]

Interesting, as said by other would be cool to see some benchmarks. Also flash attention is now supported through xformers cf #532

[0xdevalias]

I've also heard DeepSpeed MII be talked about when it comes to speed improvements of things. Unsure whether there are any things that can be understood/taken from that to improve things here:

This PR for DeepSpeed MII is a nice tutorial for SD using DeepSpeed MII and lays out some of the optimizations they do,

deepspeedai/DeepSpeed-MII#90

https://github.com/microsoft/DeepSpeed-MII/blob/7ce86af278bc4da1479c17309d87539ccf78f994/examples/benchmark/txt2img/README.md

Originally posted by @Thomas-MMJ in #532 (comment)

[Thomas-MMJ]

Here is what the krnl dev, said in response to a question someone in the reddit thread asked if they would be of benefit for accelerating yolov5,

Can I also use this project to improve inference time of projects like yolov5, etc.?

Right now the kernels cover linear layer, attention, and layer norm / rms norm. So the effect would be limited outside a transformer or assimilated. However we will increase the number of kernels, but convolution is not right now our priority

[Thomas-MMJ]

Note that there are triton implementations of convolution here, not sure how relevant they are to krnl since haven't read the code yet...

https://github.com/pytorch/pytorch/tree/c7c09722ad5ee25c5891f863e5bbd1575ad77970/torch/_inductor/triton_ops

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.