feat: add mx.fast.turboquant_attention for compressed KV cache#3340

Closed

yzamari wants to merge 8 commits into

ml-explore:mainfrom

yzamari:feature/turboquant-attention

Commits on Mar 31, 2026

feat: add mx.fast.turboquant_attention for compressed KV cache attention

Yahav Zamari
and
claude
committed
perf: apply P1 Metal shader optimizations from audit

Yahav Zamari
and
claude
committed
feat: return unnormalized (acc, m, l) for log-sum-exp merge

Yahav Zamari
and
claude
committed
feat: add 2-pass TurboQuant kernel for long sequences (N >= 1024)

Yahav Zamari
and
claude
committed
fix: 2-pass kernel batch offset for B>1

Yahav Zamari
and
claude
committed
test: add comprehensive tests for mx.fast.turboquant_attention

Yahav Zamari
and
claude
committed
style: apply pre-commit formatting + add benchmark

Yahav Zamari
and
claude
committed
feat: add 4-bit quantization support to turboquant_attention

Yahav Zamari
and
claude
committed