do you have a simple python impl? :)
Took a whole night to compute a few tokens, but I used it to do the first tests.
Also, my friend pasted the paper to claude and it produced a working basic impl instantly :D
But in all seriousness - I think MLX implementation would be doable, or a wrapper to the Metal gou functionality
This project is kind of like ultimate nerdsnipe as math is quite simple, you don’t need PhD to understand it and actually implementing things would teach you linear algebra faster vs just mindlessly doing exercises sets.
The project is a nerdsnipe for math geeks, because there are multiple small things that beg to be proven / described by math there. For example - what's the tradeoff between the number of bits we loose when embedding position vs the bits of information that we gain by knowing which bucket a weight belongs to?
In other words - is it possible that when storing weights in the bucketed form we can actually end up having a higher precision than using a regular form? For Q8 we get just 4 bits to store the weight (and 1 bit for sign, and 3 bits for location), but these 4 bits need to express numbers from a smaller range than before.
anyway here's a numerical simulation written in PyTorch for those who want to consider this algo on their projects before fully integrating: https://gist.github.com/bwasti/78d7058aad7f42dc893d906c877b7...
happy hacking