Optimize TurboQuant: O(d log d) Walsh-Hadamard Transform

[Trucker2827]

Summary

• Replace O(d²) dense rotation with O(d log d) Fast Walsh-Hadamard Transform — ~18x fewer operations for d=128, the biggest bottleneck identified in prefill/decode performance
• Replace O(d²) dense QJL projection with WHT in both _TurboQuantProdCodec and _TurboQuantPolarProdCodec
• Replace broadcasting argmin codebook search with boundary comparison, eliminating the large O(d × 2^bits) temporary tensor
• Add two Metal kernels (fast_wht_forward, fast_wht_inverse) using threadgroup shared memory for GPU-accelerated butterfly operations
• Thread unrotate_fn callback through Metal weighted-sum helpers for consistent WHT usage across quantize and decode paths
• Falls back to dense matrix for non-power-of-two dimensions (backward compatible)

Why

You mentioned in the PR description:

"This implementation is far from optimal, I'm still working on improving it to the claimed speedup results. In particular, I don't see the prefill and decode performance matching up to the claimed 8x speed up."

The #1 bottleneck is the dense random orthogonal matrix multiplication applied on every token insert AND every decode attention step. For d=128, that's 16,384 multiply-adds per vector. The Walsh-Hadamard Transform brings this to ~896 ops — an 18x reduction in the core transform. Both are theoretically valid rotations (WHT with random signs produces near-independent coordinates from the same high-d Gaussian limit as random orthogonal rotation, proven in the QuIP# literature).

Test plan

• All 15 existing tests pass
• Benchmark prefill/decode tok/s on Qwen3.5-35B-A3B with --kv-bits 3.5 --kv-quant-scheme turboquant
• Verify needle-in-a-haystack recall at 8k/32k/64k context

🤖 Generated with Claude Code

[Trucker2827]

Benchmark Results — M4 Max 128GB

Tested with our WHT optimization on Python 3.11 + MLX 0.31.1:

Model	tok/s (gen)	Peak Memory	KV Bits	Scheme
Qwen2.5-VL-7B-Instruct-4bit	85.7	5.73 GB	3.5	turboquant
Qwen3.5-35B-A3B-4bit	91.0	20.53 GB	3.5	turboquant

All 15 existing tests pass. The WHT Metal kernels are working correctly on Apple Silicon.

Would love to see a comparison on your M3 Max with the original dense rotation to quantify the speedup from WHT. Happy to iterate on any feedback!

[Blaizzy]

Could you share the results for full precision as well

And do benchmark on 8K, 32K and 64K

Basically:

• Model (preferably bf16 but quants work)
• Type
• context
• prompt tok/s
• gen tok/s
• Peak memory
• KV GB
• response correct or not / PPL

[Blaizzy]

Closing for now