Fast Hadamard Transform #1249
Fast Hadamard Transform #1249
Conversation
|
Very nice!! From an API perspective, I'm wondering if this should live in the fast namespace? The primitive does not have any transforms implemented (which is fine). I guess the question is really if we intend to implement them eventually. If yes, then maybe it makes sense to keep it as is. But if no, I would consider putting it in the fast package for now and maybe do a transformable fallback ( |
|
Agreed! The transforms are pretty simple so I added |
mlx/primitives.cpp
Outdated
| if (axes[0] == inputs[0].ndim() - 1) { | ||
| auto a = moveaxis(inputs[0], axes[0], 0, s); | ||
| auto b = hadamard_transform(a, scale_, s); | ||
| return {{moveaxis(b, 0, axes[0], s)}, axes}; |
There was a problem hiding this comment.
No need to move it back, you can just return it with 0 ie return {{b}, {0}};
mlx/primitives.h
Outdated
| DEFINE_VMAP() | ||
| DEFINE_GRADS() | ||
| DEFINE_PRINT(Hadamard) | ||
| DEFINE_DEFAULT_IS_EQUIVALENT() |
There was a problem hiding this comment.
Unfortunately the default is_equivalent is incorrect since we also need to check the scales. In general this is best left undefined unless certainly correct as it can cause quite hard to debug errors.
docs/src/python/ops.rst
Outdated
| @@ -72,6 +72,7 @@ Operations | |||
| gather_qmm | |||
| greater | |||
| greater_equal | |||
| hadamard | |||
| @@ -3950,4 +3950,37 @@ bool View::is_equivalent(const Primitive& other) const { | |||
| return (dtype_ == a_other.dtype_); | |||
| } | |||
|
|
|||
| std::pair<std::vector<array>, std::vector<int>> Hadamard::vmap( | |||
There was a problem hiding this comment.
| std::pair<std::vector<array>, std::vector<int>> Hadamard::vmap( | |
| bool Hadamard::is_equivalent(const Primitive& other) const { | |
| const Hadamard& h_other = static_cast<const Hadamard&>(other); | |
| return scale_ == h_other.scale_; | |
| } | |
| std::pair<std::vector<array>, std::vector<int>> Hadamard::vmap( |
There was a problem hiding this comment.
Also needs the declaration in primitives.h of course.
There was a problem hiding this comment.
Thanks! Added this in the latest commit.
Add a fast Hadamard transform in Metal.
Supports
n = m*2^kwherem in (1, 12, 20, 28). (e.g. Llama 3 70B has a hidden size of28672 = 28*1024).Due to shared memory limits we support
2^k <= 8192for FP32 and2^k <= 16384for FP16/BF16.Planning to use this to enable low-bit, online quantization of the KV cache similar to Quarot/SpinQuant.
Benchmarks
We get close to full bandwidth for
2^kand half bandwidth form*2^k(since we do those in two uploads). This is much faster than manually doinghadamard(N) @ xfor non-trivial batch size.