Problem
KV cache for Qwen3-TTS at 8K tokens: 28 layers × 24 heads × 8000 × 128 × 4 bytes = ~2.75 GB in f32.
Proposal
Cast K/V to float16 before cache insertion, cast back to f32 on read. MLX computes attention scores in f32 regardless of input precision, so quality impact is minimal.
Expected: ~1.4 GB KV cache (50% reduction), enabling longer generation before OOM.
Implementation
Modify mlx-rs-core/src/cache.rs KVCache::update_and_fetch() to accept an optional target dtype. Default f32, configurable to f16.
Related
- Memory module added in 7b4ef888 provides monitoring to verify the savings
- Makepad's HybridCacheTemplate (per-layer type selection) is the advanced version
Problem
KV cache for Qwen3-TTS at 8K tokens: 28 layers × 24 heads × 8000 × 128 × 4 bytes = ~2.75 GB in f32.
Proposal
Cast K/V to float16 before cache insertion, cast back to f32 on read. MLX computes attention scores in f32 regardless of input precision, so quality impact is minimal.
Expected: ~1.4 GB KV cache (50% reduction), enabling longer generation before OOM.
Implementation
Modify
mlx-rs-core/src/cache.rsKVCache::update_and_fetch()to accept an optional target dtype. Default f32, configurable to f16.Related