optimized 1x16 and 2x8 decompression/dequantization kernels now exist about aqlm HOT 8 CLOSED

Hi @BlackSamorez, I ran your new kernels on an A6000, 1 GPU and saw the same 30% speedup for 1x16 but I also got a 40% or so speedup for the 2x8. The 2x8 was already pretty blazing compared to the 1x16. Based on my timings I'm setting the breakpoints between the gemv and dequant at (the pretty tiny) batch size 6 for the 2x8 and the same for 1x16. My numbers are probably different from yours because I'm doing the scaling on the output. The 2x8 numbers are almost 70% worse when you scale the weights instead of the outputs. Here's a set of spreadsheets, check out the 2x8 and 1x16 pages. Red in the column means it's slower than doing the gemv kernel: https://docs.google.com/spreadsheets/d/1PYt90XEzUjTXJe_qVF4e3N2xnc951huamj_hPXdh_ZU/edit?usp=sharing

from aqlm.

Hi @jaemzfleming !
First of all, great work! The amount of effort put into the vllm integration is enormous just as it's practical implications!
Secondly, I have also been working on faster dequantization kernels. I've got 1x16 dequant implemented based on @efrantar 's gemv kernels. It's in the #55. I've also added your fast 1x16 dequant kernel there for comparison.
I also did a small test and I think mine is ~30% faster than that in vllm in some setups.
I'm very sorry for not bringing it up to you earlier. Hope we figure out which one is better and the amount of time spent on the redundant one wouldn't have been to large.

from aqlm.

2x8 dequant is exactly the same speed for vllm and the implementation based on the Elias's matvec code.

from aqlm.

Nice! I had hoped there could be further speed improvements. My back of envelope bandwidth calculations showed that my 1x16 was about half as fast as theoretically possible. Which is still good, but I thought maybe there would be some room left for improvement. I'll port the ones from #55 back into my PR and see if our timings agree.

With so much interest and activity there's bound to be some overlap of effort, not a problem, but a good thing, for sure.

from aqlm.

Hi @jaemzfleming I've also incorporated post matmul scaling, and it's faster indeed.
I've updated the tables and plots in #55 and I also see that 6 is around the optimal batch size to switch to matmul.

from aqlm.

Awesome! Nice to get the same numbers :) I'm now busy updating my plot in the vLLM PR with the new kernels...

from aqlm.

@BlackSamorez and I'm super happy to hear that it speeds up training as well. I had thought it would, but didn't try it. I know there were a few complaints/concerns about training times. Hopefully a moot issue now.

from aqlm.

For efficient training, the only new kernels needed are transposed matmul kernels. Using a fast dequant kernel and then transposing is already much faster than what we had before. And implementing completely separate transposed matmul/dequant would be very time consuming.

from aqlm.