🙌 ASH-KV: Asynchronous Self-Healing Cache (Zero-Latency Metal Mask Injection)

[iamrealvinnu]

Hey MLX team and community,

I’ve been building high-stakes, multi-agent reasoning loops (clinical triage, deep CoT) and hit a wall with the static nature of the KV cache. When an agent hallucinates at token 100, that logic poisons the entire forward pass for the next 8,000 tokens. Current optimizations (INT4/MLA) solve for size, but we needed to solve for truth without paying the latency penalty of full context rollbacks.

Today, I'm open-sourcing ASH-KV (Asynchronous Self-Healing Cache).

It acts as an "Active Immune System" for MLX inference by exploiting Apple Silicon's Unified Memory architecture to run a dual-threaded manifold:

1. Primary Thread (GPU): Runs the standard mlx_lm autoregressive loop at maximum throughput.
2. Ghost Critic (Daemon): A parallel, asynchronous process (e.g., a PRM or ANE-compiled validator) that monitors the shared memory manifold for logical drift.
3. The Exploit: When the Critic flags a hallucinated node, ASH-KV doesn't reallocate physical memory. Instead, it triggers a surgically fused Metal kernel (via @mx.compile) that natively injects a -10000.0 penalty into the attention mask at the exact poisoned index.

The primary model "forgets" the bad logic on the very next token pass with absolute zero memory reallocation overhead.

The M4 Telemetry:

• Critic Overhead: 0.00 ms (Fully Asynchronous)
• Mutation Latency: < 0.1 ms (Fused Metal Kernel)
• Throughput: 100% maintained during healing events

Massive respect to @awni and the core team, the @mx.compile decorator made fusing this surgical strike incredibly elegant.

Links:
💻 Source & Terminal UI: iamrealvinnu/mlx-ash-kv
🔬 Live Attention Visualizer: Hugging Face Space

Would love any feedback from the core team on how we are handling the concurrent unified memory reads/writes, or if there are further low-level Metal optimizations we can squeeze out of the mask injection!