Add vector-search-tool: MLX TurboQuant vector search example

.github/workflows/pull_request.yml

@@ -109,3 +109,4 @@ jobs:
           xcodebuild -scheme embedder-tool -skipMacroValidation
           xcodebuild -scheme image-tool -skipMacroValidation
           xcodebuild -scheme mnist-tool -skipMacroValidation
+          xcodebuild -scheme vector-search-tool -skipMacroValidation

ACKNOWLEDGMENTS.md

@@ -7,6 +7,7 @@
 MLX Swift was developed with contributions from the following individuals:
 
 - [John Mai](https://github.com/johnmai-dev): Added support for multiple models (Qwen2, Starcoder2, InternLM2, Qwen3, Qwen3 MoE, GLM-4, MiMo, BitNet, SmolLM3, LFM2, Baichuan-M1).
+- [Joel Nishanth](https://offlyn.ai): Added `vector-search-tool` demonstrating MLX-accelerated TurboQuant vector search via [mlx-turbovec-swift](https://github.com/offlyn-ai/mlx-turbovec-swift).
 
 <a href="https://github.com/ml-explore/mlx-swift-examples/graphs/contributors">
   <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/mlx-swift-examples&anon=0&columns=20&max=100&r=true" />

Package.swift

@@ -16,6 +16,8 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/ml-explore/mlx-swift", .upToNextMinor(from: "0.31.4")),
+        .package(url: "https://github.com/offlyn-ai/mlx-turbovec-swift.git", from: "0.1.0"),
+        .package(url: "https://github.com/apple/swift-argument-parser", from: "1.4.0"),
 
         // Note: used by StableDiffusion library to download weights
         .package(
@@ -59,6 +61,17 @@ let package = Package(
                 .enableExperimentalFeature("StrictConcurrency")
             ]
         ),
+        .executableTarget(
+            name: "vector-search-tool",
+            dependencies: [
+                .product(name: "TurboVec", package: "mlx-turbovec-swift"),
+                .product(name: "ArgumentParser", package: "swift-argument-parser"),
+            ],
+            path: "Tools/vector-search-tool",
+            exclude: [
+                "README.md"
+            ]
+        ),
     ]
 )
 

README.md

@@ -35,6 +35,10 @@ examples use models implemented in [MLX Swift LM](https://github.com/ml-explore/
 - [mnist-tool](Tools/mnist-tool/README.md): A command line tool for training a
   a LeNet on MNIST.
 
+- [vector-search-tool](Tools/vector-search-tool/README.md): A command line tool
+  demonstrating MLX-accelerated TurboQuant vector search with compression and
+  recall benchmarks. Uses [mlx-turbovec-swift](https://github.com/offlyn-ai/mlx-turbovec-swift).
+
 ## Numerical Computing
 
 Examples that use MLX for general numerical computing (no ML model involved),

Tools/vector-search-tool/README.md

@@ -0,0 +1,80 @@
+# vector-search-tool
+
+A command-line tool demonstrating [mlx-turbovec-swift](https://github.com/offlyn-ai/mlx-turbovec-swift) — TurboQuant vector quantization with optional MLX GPU acceleration.
+
+Unlike `embedder-tool` (which embeds text with MLX Embedders and stores raw vectors as JSON), this tool shows **compressed vector search**: 768-dim embeddings shrink from 3,072 bytes to 384 bytes (4-bit) while maintaining >90% recall@10.
+
+### Building
+
+Build the `vector-search-tool` scheme in Xcode, or use Swift Package Manager:
+
+```bash
+swift build --target vector-search-tool
+```
+
+### Running: Command Line
+
+Use the `mlx-run` helper after building in Xcode:
+
+```bash
+./mlx-run vector-search-tool --gpu --dim 768 --count 10000 --queries 100 --k 10
+```
+
+Pass `--debug` after `mlx-run` to run the Debug configuration.
+
+Enable MLX GPU acceleration with `--gpu`. Without it, TurboVec falls back to Accelerate CPU paths (useful for CI and environments without Metal).
+
+Write a structured JSON report for LLM analysis:
+
+```bash
+./mlx-run vector-search-tool --gpu --json /tmp/turbovec-report.json
+```
+
+### Running: Xcode
+
+Configure scheme arguments (Product > Scheme > Edit Scheme > Run > Arguments):
+
+```
+--gpu --dim 768 --count 10000 --queries 50 --k 10
+```
+
+Then press <kbd>⌘</kbd>+<kbd>R</kbd> to run.
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--dim` | 768 | Vector dimension (multiple of 8) |
+| `--count` | 10000 | Vectors to index |
+| `--queries` | 100 | Search queries to run |
+| `--k` | 10 | Top-k neighbors |
+| `--bits` | 4 | Quantization bit width (2, 3, or 4) |
+| `--gpu` | off | Enable MLX GPU acceleration |
+| `--json` | — | Save JSON benchmark report to path |
+
+### Expected Output
+
+```
+▸ MLX GPU acceleration enabled
+▸ Generating 10000 random unit vectors (d=768)...
+▸ Indexing with 4-bit TurboQuant...
+  ✓ Indexed in 842.3ms
+  ✓ Compression: 8.0× (30720000 → 3840000 bytes)
+▸ Running 100 searches (k=10)...
+  ✓ Mean latency: 0.412ms
+  ✓ P99 latency:  0.891ms
+  ✓ QPS:          2427
+▸ Measuring recall vs brute-force baseline...
+  ✓ Recall@1:  0.9100
+  ✓ Recall@10: 0.9450
+  ✓ Brute-force mean: 12.340ms
+  ✓ Speedup: 29.9×
+
+Summary: 8.0× compression, R@1=0.9100, 29.9× faster than brute-force
+```
+
+See also:
+
+- [mlx-turbovec-swift](https://github.com/offlyn-ai/mlx-turbovec-swift) — the TurboQuant library
+- [embedder-tool](../embedder-tool/README.md) — MLX Embedders for text embedding + JSON index
+- [MLX troubleshooting](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/troubleshooting)

Tools/vector-search-tool/VectorSearchTool.swift

@@ -0,0 +1,225 @@
+import ArgumentParser
+import Foundation
+import TurboVec
+
+/// Command-line demo of MLX-accelerated TurboQuant vector search.
+///
+/// Generates synthetic unit vectors, indexes them with TurboVec, runs top-k
+/// search, and reports compression, latency, and recall vs brute-force cosine.
+@main
+struct VectorSearchTool: ParsableCommand {
+    static let configuration = CommandConfiguration(
+        abstract: "Demonstrate MLX-accelerated TurboQuant vector search",
+        discussion: """
+            Builds a TurboQuant index over random vectors and compares approximate \
+            search against an exact brute-force baseline. Enable MLX GPU acceleration \
+            before indexing for faster batch rotation.
+            """
+    )
+
+    @Option(name: .long, help: "Vector dimension (multiple of 8)")
+    var dim: Int = 768
+
+    @Option(name: .long, help: "Number of vectors to index")
+    var count: Int = 10_000
+
+    @Option(name: .long, help: "Number of search queries")
+    var queries: Int = 100
+
+    @Option(name: .long, help: "Top-k neighbors to retrieve")
+    var k: Int = 10
+
+    @Option(name: .long, help: "Quantization bit width (2, 3, or 4)")
+    var bits: Int = 4
+
+    @Flag(name: .long, help: "Enable MLX GPU acceleration via TurboVec")
+    var gpu = false
+
+    @Option(name: .long, help: "Write structured JSON benchmark report to path")
+    var json: String?
+
+    mutating func run() throws {
+        guard dim > 0, dim % 8 == 0 else {
+            throw ValidationError("--dim must be a positive multiple of 8")
+        }
+        guard let bitWidth = BitWidth(rawValue: UInt8(bits)) else {
+            throw ValidationError("--bits must be 2, 3, or 4")
+        }
+        guard count > 0, queries > 0, k > 0 else {
+            throw ValidationError("--count, --queries, and --k must be positive")
+        }
+
+        if gpu {
+            MLXBackend.enableGPU()
+            print("▸ MLX GPU acceleration enabled")
+        } else {
+            print("▸ CPU-only mode (pass --gpu to enable MLX acceleration)")
+        }
+
+        print("▸ Generating \(count) random unit vectors (d=\(dim))...")
+        var rng = SplitMix64(seed: 42)
+        let vectors = Self.generateVectors(n: count, dim: dim, rng: &rng)
+        let queryVectors = Self.generateVectors(n: queries, dim: dim, rng: &rng)
+
+        print("▸ Indexing with \(bits)-bit TurboQuant...")
+        let index = TurboQuantIndex(dim: dim, bitWidth: bitWidth)
+        let indexStart = CFAbsoluteTimeGetCurrent()
+        try index.add(vectors)
+        let indexMs = (CFAbsoluteTimeGetCurrent() - indexStart) * 1000
+
+        let rawBytes = count * dim * MemoryLayout<Float>.size
+        let indexBytes = index.indexSizeBytes
+        let compression = Double(rawBytes) / Double(indexBytes)
+
+        print("  ✓ Indexed in \(String(format: "%.1f", indexMs))ms")
+        print("  ✓ Compression: \(String(format: "%.1f", compression))× (\(rawBytes) → \(indexBytes) bytes)")
+
+        print("▸ Running \(queries) searches (k=\(k))...")
+        var latencies = [Double]()
+        var approximateResults = [[SearchHit]]()
+
+        for query in queryVectors {
+            let start = CFAbsoluteTimeGetCurrent()
+            let hits = try index.search(query: query, k: k)
+            latencies.append((CFAbsoluteTimeGetCurrent() - start) * 1000)
+            approximateResults.append(hits)
+        }
+
+        let searchStats = Self.latencyStats(latencies)
+        print("  ✓ Mean latency: \(String(format: "%.3f", searchStats.meanMs))ms")
+        print("  ✓ P99 latency:  \(String(format: "%.3f", searchStats.p99Ms))ms")
+        print("  ✓ QPS:          \(String(format: "%.0f", 1000.0 / searchStats.meanMs))")
+
+        print("▸ Measuring recall vs brute-force baseline...")
+        var recall1Sum = 0.0
+        var recallKSum = 0.0
+        var bfLatencies = [Double]()
+
+        for (qi, query) in queryVectors.enumerated() {
+            let start = CFAbsoluteTimeGetCurrent()
+            let exact = BruteForceSearch.search(query: query, vectors: vectors, k: k)
+            bfLatencies.append((CFAbsoluteTimeGetCurrent() - start) * 1000)
+
+            let approx = approximateResults[qi]
+            let exactTopK = Set(exact.prefix(k).map(\.index))
+            let approxIndices = approx.prefix(k).map(\.index)
+
+            if !exact.isEmpty, !approxIndices.isEmpty, approxIndices[0] == exact[0].index {
+                recall1Sum += 1
+            }
+            let overlap = approxIndices.filter { exactTopK.contains($0) }.count
+            recallKSum += Double(overlap) / Double(min(k, exact.count))
+        }
+
+        let recall1 = recall1Sum / Double(queries)
+        let recallK = recallKSum / Double(queries)
+        let bfStats = Self.latencyStats(bfLatencies)
+        let speedup = bfStats.meanMs / searchStats.meanMs
+
+        print("  ✓ Recall@1:  \(String(format: "%.4f", recall1))")
+        print("  ✓ Recall@\(k): \(String(format: "%.4f", recallK))")
+        print("  ✓ Brute-force mean: \(String(format: "%.3f", bfStats.meanMs))ms")
+        print("  ✓ Speedup: \(String(format: "%.1f", speedup))×")
+
+        let report = BenchmarkReport(
+            hardware: Self.detectHardware(),
+            config: BenchmarkReport.TestConfig(
+                dimensions: dim,
+                bitWidth: bits,
+                numVectors: count,
+                numQueries: queries,
+                k: k
+            ),
+            search: BenchmarkReport.SearchResult(
+                latency: searchStats,
+                queriesPerSecond: 1000.0 / searchStats.meanMs,
+                recallAt1: recall1,
+                recallAtK: recallK
+            ),
+            memory: BenchmarkReport.MemoryUsage(
+                rawBytes: rawBytes,
+                indexBytes: indexBytes,
+                compressionRatio: compression,
+                bytesPerVector: Double(indexBytes) / Double(count)
+            ),
+            indexing: BenchmarkReport.IndexingResult(
+                totalMs: indexMs,
+                perVectorUs: indexMs * 1000.0 / Double(count),
+                vectorsPerSecond: Double(count) / (indexMs / 1000.0)
+            ),
+            bruteForceBaseline: BenchmarkReport.SearchResult(
+                latency: bfStats,
+                queriesPerSecond: 1000.0 / bfStats.meanMs,
+                recallAt1: 1.0,
+                recallAtK: 1.0
+            ),
+            timestamp: ISO8601DateFormatter().string(from: Date()),
+            version: "0.1.0"
+        )
+
+        print()
+        print("Summary: \(String(format: "%.1f", compression))× compression, " +
+            "R@1=\(String(format: "%.4f", recall1)), " +
+            "\(String(format: "%.1f", speedup))× faster than brute-force")
+
+        let jsonText = try report.jsonString()
+        if let json {
+            try jsonText.write(toFile: json, atomically: true, encoding: .utf8)
+            print("Report saved to \(json)")
+        }
+
+        print()
+        print("--- JSON Report ---")
+        print(jsonText)
+    }
+
+    private static func generateVectors(n: Int, dim: Int, rng: inout SplitMix64) -> [[Float]] {
+        (0..<n).map { _ in
+            var vector = (0..<dim).map { _ -> Float in
+                let u1 = max(Float(rng.nextUniform()), Float.leastNormalMagnitude)
+                let u2 = Float(rng.nextUniform())
+                return sqrtf(-2.0 * logf(u1)) * cosf(2.0 * .pi * u2)
+            }
+            VectorMath.normalize(&vector)
+            return vector
+        }
+    }
+
+    private static func latencyStats(_ values: [Double]) -> BenchmarkReport.LatencyStats {
+        let sorted = values.sorted()
+        let n = sorted.count
+        return BenchmarkReport.LatencyStats(
+            p50Ms: sorted[n / 2],
+            p95Ms: sorted[Int(Double(n) * 0.95)],
+            p99Ms: sorted[Int(Double(n) * 0.99)],
+            meanMs: sorted.reduce(0, +) / Double(n),
+            minMs: sorted.first ?? 0,
+            maxMs: sorted.last ?? 0
+        )
+    }
+
+    private static func detectHardware() -> BenchmarkReport.HardwareInfo {
+        var chip = "Unknown"
+        #if arch(arm64)
+        chip = sysctlString("machdep.cpu.brand_string") ?? "Apple Silicon"
+        #else
+        chip = sysctlString("machdep.cpu.brand_string") ?? "x86_64"
+        #endif
+
+        return BenchmarkReport.HardwareInfo(
+            chip: chip,
+            cores: ProcessInfo.processInfo.activeProcessorCount,
+            memoryGB: Int(ProcessInfo.processInfo.physicalMemory / (1024 * 1024 * 1024)),
+            os: ProcessInfo.processInfo.operatingSystemVersionString
+        )
+    }
+
+    private static func sysctlString(_ name: String) -> String? {
+        var size = 0
+        sysctlbyname(name, nil, &size, nil, 0)
+        guard size > 0 else { return nil }
+        var buffer = [CChar](repeating: 0, count: size)
+        sysctlbyname(name, &buffer, &size, nil, 0)
+        return String(cString: buffer)
+    }
+}