[perf] Cache MLX H3 text refinement and speed up attention (#2002)

Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>

docs/inference/fasth3_omniref_mlx.md

@@ -73,6 +73,34 @@ Quantized checkpoints pick a linear path per layer. `--quant-gemm` (or
 
 The same flag is on the example: `--quant-gemm {auto,dequant,qmm}`.
 
+### Opt-in M5 attention backend
+
+`--vsa-impl sdpa` evaluates the trained VSA tile mask through fused MLX SDPA,
+sharing each attention head's K/V across its query tiles. On M5, the built-in
+matrix operations can be faster than the custom SIMD kernel even though they
+visit masked keys. The default `auto` selection is unchanged.
+
+This backend keeps the reference resolution, selected tiles, ragged-row masks,
+and compression gate. Accumulation differs from the SIMD backend, so outputs
+are not bit-identical. It supports tiles 64/128 with head dimension 128 and
+raises if fused SDPA is unavailable. Prefix K/V reuse is not supported; do not
+combine it with `--reuse-prefix-kv`. Query tiles are processed in chunks of at
+most 128, capped at 256 MiB per boolean-mask chunk with two chunks in flight.
+
+`--vsa-impl nax` is a separate opt-in backend for Apple M5 GPUs. It uses Metal
+TensorOps to visit only the selected key blocks and their valid rows, avoiding
+the dense mask and masked-key traversal of the SDPA backend. Dense prefix
+queries and the trained compression branch keep their existing implementations.
+The selected reference tokens, output resolution and denoising steps are unchanged.
+
+The NAX backend uses private helper headers from the installed MLX package;
+it accepts only the header hashes validated with MLX 0.32.4. Changed headers need
+numerical and generation validation before their hashes are accepted. Missing headers, unsupported GPUs
+and compilation failures raise errors rather than silently selecting another
+backend. Tiles 64/128 and head dimension 128 are supported. Prefix K/V reuse is
+not supported. Floating-point accumulation differs from both SIMD and SDPA,
+so the generated clip can differ even with the same seed. `auto` remains unchanged.
+
 ## Generate
 
 ```bash