Top-K Identification on 1M 128-dimensional vectors (99% recall)
6.55Matmul Latency (ms)Generalized Two-Stage Approximate Top-K
Evaluation results
| Method | Links | ||||
|---|---|---|---|---|---|
| Generalized Two-Stage Approximate Top-KK'=4, Fusion=fused2025.06 | 6.55 | — | 3.51 | 10 | |
| Generalized Two-Stage Approximate Top-KK'=4, Fusion=unfused2025.06 | 7.31 | 10.8 | 3.51 | 22 | |
| jax.lax.top_kimplementation=exact Top-K2025.06 | 7.32 | — | 587 | 594 | |
| Generalized Two-Stage Approximate Top-KK'=1, Fusion=unfused2025.06 | 7.32 | 6.58 | 50 | 64 | |
| Generalized Two-Stage Approximate Top-KK'=1, Fusion=fused2025.06 | 9.03 | — | 50 | 59 | |
| jax.lax.approx_max_kK'=1, Fusion=fused2025.06 | 9.06 | — | 118 | 127 |