faiss 0.6.1 → 0.6.3

data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp

@@ -11,8 +11,10 @@
 #include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/test/TestUtils.h>
 #include <faiss/gpu/utils/DeviceUtils.h>
+#include <faiss/impl/IDSelector.h>
 #include <gtest/gtest.h>
 #include <cmath>
+#include <limits> // IWYU pragma: keep
 #include <sstream>
 #include <vector>
 
@@ -74,6 +76,7 @@ void runCopyToTest(faiss::ScalarQuantizer::QuantizerType qtype) {
 
     auto config = GpuIndexIVFScalarQuantizerConfig();
     config.device = opt.device;
+    config.use_cuvs = false;
 
     GpuIndexIVFScalarQuantizer gpuIndex(
             &res, opt.dim, opt.numCentroids, qtype, METRIC_L2, true, config);
@@ -163,6 +166,7 @@ void runCopyFromTest(faiss::ScalarQuantizer::QuantizerType qtype) {
 
     auto config = GpuIndexIVFScalarQuantizerConfig();
     config.device = opt.device;
+    config.use_cuvs = false;
 
     GpuIndexIVFScalarQuantizer gpuIndex(
             &res,
@@ -223,6 +227,182 @@ TEST(TestGpuIndexIVFScalarQuantizer, CopyFrom_4bit_uniform) {
     runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_4bit_uniform);
 }
 
+#if defined USE_NVIDIA_CUVS
+void runCuvsSQ8Test(faiss::MetricType metricType) {
+    using namespace faiss;
+    using namespace faiss::gpu;
+
+    Options opt;
+    opt.indicesOpt = INDICES_64_BIT;
+
+    std::vector<float> trainVecs = randVecs(opt.numTrain, opt.dim);
+    std::vector<float> addVecs = randVecs(opt.numAdd, opt.dim);
+    std::vector<float> queryVecs = randVecs(opt.numQuery, opt.dim);
+
+    StandardGpuResources res;
+    res.noTempMemory();
+
+    auto config = GpuIndexIVFScalarQuantizerConfig();
+    config.device = opt.device;
+    config.indicesOptions = INDICES_64_BIT;
+    config.use_cuvs = true;
+    if (!should_use_cuvs(config)) {
+        GTEST_SKIP() << "cuVS is not enabled for this GPU";
+    }
+
+    GpuIndexIVFScalarQuantizer gpuIndex(
+            &res,
+            opt.dim,
+            opt.numCentroids,
+            ScalarQuantizer::QT_8bit,
+            metricType,
+            true,
+            config);
+    gpuIndex.train(opt.numTrain, trainVecs.data());
+    gpuIndex.add(opt.numAdd, addVecs.data());
+    gpuIndex.nprobe = opt.nprobe;
+
+    std::vector<float> distances(opt.numQuery * opt.k);
+    std::vector<idx_t> labels(opt.numQuery * opt.k);
+    gpuIndex.search(
+            opt.numQuery,
+            queryVecs.data(),
+            opt.k,
+            distances.data(),
+            labels.data());
+    EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
+    EXPECT_NE(labels[0], idx_t(-1));
+
+    IndexFlatL2 quantizerL2(opt.dim);
+    IndexFlatIP quantizerIP(opt.dim);
+    Index* quantizer = metricType == METRIC_L2 ? (Index*)&quantizerL2
+                                               : (Index*)&quantizerIP;
+
+    IndexIVFScalarQuantizer cpuCopy(
+            quantizer,
+            opt.dim,
+            opt.numCentroids,
+            ScalarQuantizer::QT_8bit,
+            metricType,
+            true);
+    gpuIndex.copyTo(&cpuCopy);
+    cpuCopy.nprobe = opt.nprobe;
+
+    EXPECT_EQ(cpuCopy.ntotal, gpuIndex.ntotal);
+    EXPECT_EQ(cpuCopy.d, gpuIndex.d);
+    EXPECT_EQ(cpuCopy.nlist, gpuIndex.getNumLists());
+    EXPECT_EQ(cpuCopy.sq.qtype, ScalarQuantizer::QT_8bit);
+    EXPECT_EQ(cpuCopy.sq.trained.size(), 2 * opt.dim);
+    EXPECT_TRUE(cpuCopy.by_residual);
+
+    compareIndices(
+            cpuCopy,
+            gpuIndex,
+            opt.numQuery,
+            opt.dim,
+            opt.k,
+            opt.toString(),
+            0.08f,
+            0.2f,
+            0.05f);
+
+    GpuIndexIVFScalarQuantizer copiedGpuIndex(
+            &res, 1, 1, ScalarQuantizer::QT_4bit, METRIC_L2, false, config);
+    copiedGpuIndex.copyFrom(&cpuCopy);
+    copiedGpuIndex.nprobe = opt.nprobe;
+    EXPECT_EQ(copiedGpuIndex.ntotal, cpuCopy.ntotal);
+    testIVFEquality(cpuCopy, copiedGpuIndex);
+
+    compareIndices(
+            cpuCopy,
+            copiedGpuIndex,
+            opt.numQuery,
+            opt.dim,
+            opt.k,
+            opt.toString(),
+            0.08f,
+            0.2f,
+            0.05f);
+
+    SearchParametersIVF searchParams;
+    searchParams.nprobe = opt.nprobe;
+    TestIDSelectorStruct selectorStruct(opt.numAdd);
+    for (auto& [selectorName, selector] : selectorStruct.selector_map) {
+        searchParams.sel = selector.get();
+        testIDSelectorSearch(
+                &gpuIndex,
+                &searchParams,
+                queryVecs,
+                opt.numQuery,
+                opt.k,
+                selectorName);
+    }
+
+    GpuIndexIVFScalarQuantizer shiftedGpuIndex(
+            &res,
+            opt.dim,
+            opt.numCentroids,
+            ScalarQuantizer::QT_8bit,
+            metricType,
+            true,
+            config);
+    shiftedGpuIndex.train(opt.numTrain, trainVecs.data());
+    std::vector<idx_t> shiftedIds(opt.numAdd);
+    idx_t shiftedIdBase = opt.numAdd * 4 + 123;
+    for (int i = 0; i < opt.numAdd; ++i) {
+        shiftedIds[i] = shiftedIdBase + i;
+    }
+    shiftedGpuIndex.add_with_ids(opt.numAdd, addVecs.data(), shiftedIds.data());
+    shiftedGpuIndex.nprobe = opt.nprobe;
+
+    SearchParametersIVF shiftedSearchParams;
+    shiftedSearchParams.nprobe = opt.nprobe;
+    IDSelectorRange shiftedSelector(shiftedIdBase, shiftedIdBase + opt.numAdd);
+    shiftedSearchParams.sel = &shiftedSelector;
+    std::vector<float> shiftedDistances(opt.numQuery * opt.k);
+    std::vector<idx_t> shiftedLabels(opt.numQuery * opt.k, -1);
+    shiftedGpuIndex.search(
+            opt.numQuery,
+            queryVecs.data(),
+            opt.k,
+            shiftedDistances.data(),
+            shiftedLabels.data(),
+            &shiftedSearchParams);
+
+    bool foundShiftedLabel = false;
+    for (idx_t label : shiftedLabels) {
+        if (label >= 0) {
+            foundShiftedLabel = true;
+            EXPECT_TRUE(shiftedSelector.is_member(label));
+        }
+    }
+    EXPECT_TRUE(foundShiftedLabel);
+
+    auto nanQueries = queryVecs;
+    nanQueries[opt.dim + 3] = std::numeric_limits<float>::quiet_NaN();
+    gpuIndex.search(
+            2, nanQueries.data(), opt.k, distances.data(), labels.data());
+    for (int j = 0; j < opt.k; ++j) {
+        EXPECT_EQ(labels[opt.k + j], idx_t(-1));
+        EXPECT_EQ(distances[opt.k + j], std::numeric_limits<float>::max());
+    }
+
+    gpuIndex.reset();
+    EXPECT_EQ(gpuIndex.ntotal, 0);
+    for (int i = 0; i < opt.numCentroids; ++i) {
+        EXPECT_EQ(gpuIndex.getListLength(i), 0);
+    }
+}
+
+TEST(TestCuvsGpuIndexIVFScalarQuantizer, SQ8_L2) {
+    runCuvsSQ8Test(faiss::METRIC_L2);
+}
+
+TEST(TestCuvsGpuIndexIVFScalarQuantizer, SQ8_IP) {
+    runCuvsSQ8Test(faiss::METRIC_INNER_PRODUCT);
+}
+#endif
+
 int main(int argc, char** argv) {
     testing::InitGoogleTest(&argc, argv);
 

data/vendor/faiss/faiss/gpu_metal/MetalDistance.h

@@ -0,0 +1,87 @@
+// @lint-ignore-every LICENSELINT
+/**
+ * Copyright (c) Meta Platforms, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * IVF distance computation and scan dispatch for Metal backend.
+ */
+
+#pragma once
+
+#import <Metal/Metal.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+namespace faiss {
+namespace gpu_metal {
+
+class MetalResources;
+
+int getMetalDistanceMaxK();
+
+bool runMetalComputeNorms(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> vectors,
+        int nb,
+        int d,
+        id<MTLBuffer> normsBuf,
+        bool waitForCompletion = true);
+
+bool runMetalIVFFlatScan(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> queries,
+        id<MTLBuffer> codes,
+        id<MTLBuffer> ids,
+        id<MTLBuffer> listOffset,
+        id<MTLBuffer> listLength,
+        id<MTLBuffer> coarseAssign,
+        int nq,
+        int d,
+        int k,
+        int nprobe,
+        bool isL2,
+        id<MTLBuffer> outDistances,
+        id<MTLBuffer> outIndices,
+        id<MTLBuffer> perListDistBuf,
+        id<MTLBuffer> perListIdxBuf,
+        id<MTLBuffer> interleavedCodes = nil,
+        id<MTLBuffer> interleavedCodesOffset = nil,
+        bool waitForCompletion = true);
+
+bool runMetalIVFFlatFullSearch(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> queries,
+        int nq,
+        int d,
+        int k,
+        int nprobe,
+        bool isL2,
+        id<MTLBuffer> centroids,
+        int nlist,
+        id<MTLBuffer> codes,
+        id<MTLBuffer> ids,
+        id<MTLBuffer> listOffset,
+        id<MTLBuffer> listLength,
+        id<MTLBuffer> outDistances,
+        id<MTLBuffer> outIndices,
+        id<MTLBuffer> perListDistBuf,
+        id<MTLBuffer> perListIdxBuf,
+        id<MTLBuffer> coarseDistBuf,
+        id<MTLBuffer> coarseIdxBuf,
+        id<MTLBuffer> distMatrixBuf,
+        id<MTLBuffer> centroidNormsBuf = nil,
+        int avgListLen = 256,
+        id<MTLBuffer> interleavedCodes = nil,
+        id<MTLBuffer> interleavedCodesOffset = nil,
+        bool centroidsAreFP16 = false,
+        bool waitForCompletion = true);
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/MetalIndex.h

@@ -11,6 +11,7 @@
 #pragma once
 
 #include <faiss/Index.h>
+#include <faiss/gpu/GpuIndicesOptions.h>
 #include <faiss/gpu_metal/MetalResources.h>
 #include <memory>
 
@@ -20,6 +21,12 @@ namespace gpu_metal {
 /// Configuration for Metal index (mirrors GpuIndexConfig roles).
 struct MetalIndexConfig {
     int device = 0;
+
+    bool useFloat16CoarseQuantizer = false;
+
+    faiss::gpu::IndicesOptions indicesOptions = faiss::gpu::INDICES_64_BIT;
+
+    bool interleavedLayout = true;
 };
 
 /// Base class for Metal-backed indexes. Mirrors faiss::gpu::GpuIndex.

data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFFlat.h

@@ -0,0 +1,177 @@
+// @lint-ignore-every LICENSELINT
+/**
+ * Copyright (c) Meta Platforms, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * Minimal Metal IVFFlat wrapper.
+ *
+ */
+
+#pragma once
+
+#import <Metal/Metal.h>
+
+#include <faiss/IndexIVFFlat.h>
+#include <faiss/gpu/GpuIndicesOptions.h>
+#include <faiss/gpu_metal/MetalIndex.h>
+
+#include <memory>
+
+namespace faiss {
+namespace gpu_metal {
+
+class MetalIVFFlatImpl;
+
+/// IVFFlat index wrapper for Metal backend.
+/// Currently delegates to an internal CPU IndexIVFFlat; later phases
+/// may move list scanning to GPU.
+class MetalIndexIVFFlat : public MetalIndex {
+   public:
+    struct AppendDebugStats {
+        size_t relayoutEvents = 0;
+        size_t movedLists = 0;
+        size_t movedVectors = 0;
+        size_t reusedSegmentAllocs = 0;
+        size_t tailSegmentAllocs = 0;
+        size_t reusedCapacityVecs = 0;
+        size_t tailCapacityVecs = 0;
+        size_t tailShrinkEvents = 0;
+        size_t tailShrunkVecs = 0;
+    };
+
+    /// Construct empty IVFFlat index with its own CPU quantizer.
+    MetalIndexIVFFlat(
+            std::shared_ptr<MetalResources> resources,
+            int dims,
+            idx_t nlist,
+            faiss::MetricType metric,
+            float metricArg = 0.0f,
+            MetalIndexConfig config = MetalIndexConfig());
+
+    /// Construct empty IVFFlat index with caller-provided coarse quantizer.
+    /// If ownFields is true, this index takes ownership of `coarseQuantizer`.
+    MetalIndexIVFFlat(
+            std::shared_ptr<MetalResources> resources,
+            faiss::Index* coarseQuantizer,
+            int dims,
+            idx_t nlist,
+            faiss::MetricType metric,
+            float metricArg = 0.0f,
+            MetalIndexConfig config = MetalIndexConfig(),
+            bool ownFields = false);
+
+    /// Construct from an existing CPU IndexIVFFlat (used by cloners later).
+    MetalIndexIVFFlat(
+            std::shared_ptr<MetalResources> resources,
+            const faiss::IndexIVFFlat* cpuIndex,
+            MetalIndexConfig config = MetalIndexConfig());
+
+    ~MetalIndexIVFFlat() override;
+
+    void train(idx_t n, const float* x) override;
+    void add(idx_t n, const float* x) override;
+    void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
+    void reset() override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    /// Search with caller-provided coarse assignments (skips coarse quantizer).
+    /// @param assign      Coarse list assignments (n x nprobe), row-major idx_t
+    /// @param centroid_dis Distances to assigned centroids (n x nprobe); unused
+    ///                     by GPU scan but accepted for API compatibility
+    /// @param store_pairs  Ignored (always false for GPU path)
+    void search_preassigned(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            const idx_t* assign,
+            const float* centroid_dis,
+            float* distances,
+            idx_t* labels,
+            bool store_pairs,
+            const IVFSearchParameters* params = nullptr,
+            IndexIVFStats* stats = nullptr) const;
+
+    /// Copy from a CPU IndexIVFFlat (helper for future cloner support).
+    void copyFrom(const faiss::IndexIVFFlat* index);
+
+    /// Copy to a CPU IndexIVFFlat.
+    void copyTo(faiss::IndexIVFFlat* index) const;
+
+    /// Reconstruct a single stored vector by internal key.
+    void reconstruct(idx_t key, float* recons) const override;
+
+    /// Reconstruct n contiguous stored vectors starting at i0.
+    void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
+
+    /// Re-upload coarse quantizer centroids to GPU after external changes.
+    void updateQuantizer();
+
+    /// Return the vector indices in inverted list `listId`.
+    std::vector<idx_t> getListIndices(idx_t listId) const;
+
+    /// Return raw vector data from inverted list `listId`.
+    std::vector<float> getListVectorData(idx_t listId) const;
+
+    /// Release unused GPU memory.
+    void reclaimMemory();
+
+    /// Pre-allocate GPU storage for the given total number of vectors.
+    void reserveMemory(idx_t numVecs);
+
+    /// Accessors (needed by cloner and tests).
+    idx_t nlist() const;
+    size_t nprobe() const;
+    bool interleavedLayout() const;
+    faiss::gpu::IndicesOptions indicesOptions() const;
+    AppendDebugStats appendDebugStats() const;
+    void resetAppendDebugStats();
+
+   private:
+    std::unique_ptr<faiss::IndexIVFFlat> cpuIndex_;
+    std::unique_ptr<MetalIVFFlatImpl> gpuIvf_;
+    faiss::gpu::IndicesOptions indicesOptions_;
+    bool interleavedLayout_;
+
+    // Persistent search buffers — allocated once, grown lazily.
+    // Declared mutable so search() (const) can resize them.
+    mutable id<MTLBuffer> searchQueriesBuf_ = nil;
+    mutable id<MTLBuffer> searchCoarseBuf_ = nil;
+    mutable id<MTLBuffer> searchOutDistBuf_ = nil;
+    mutable id<MTLBuffer> searchOutIdxBuf_ = nil;
+    mutable size_t searchQueriesCap_ = 0; // bytes
+    mutable size_t searchCoarseCap_ = 0;
+    mutable size_t searchOutDistCap_ = 0;
+    mutable size_t searchOutIdxCap_ = 0;
+    mutable id<MTLBuffer> searchPerListDistBuf_ = nil;
+    mutable id<MTLBuffer> searchPerListIdxBuf_ = nil;
+    mutable size_t searchPerListDistCap_ = 0;
+    mutable size_t searchPerListIdxCap_ = 0;
+
+    // GPU coarse quantizer buffers (cached, rebuilt on train)
+    mutable id<MTLBuffer> centroidBuf_ = nil;
+    mutable id<MTLBuffer> centroidNormsBuf_ = nil; // pre-computed ||c||²
+    mutable id<MTLBuffer> coarseOutDistBuf_ = nil;
+    mutable id<MTLBuffer> coarseOutIdxBuf_ = nil;
+    mutable size_t coarseOutDistCap_ = 0;
+    mutable size_t coarseOutIdxCap_ = 0;
+    mutable id<MTLBuffer> distMatrixBuf_ = nil;
+    mutable size_t distMatrixCap_ = 0;
+
+    /// Ensures buf is at least `needed` bytes, reallocating if necessary.
+    void ensureSearchBuf_(id<MTLBuffer>& buf, size_t& cap, size_t needed) const;
+
+    /// (Re)uploads quantizer centroids to centroidBuf_.
+    void uploadCentroids_() const;
+};
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFPQ.h

@@ -0,0 +1,88 @@
+// @lint-ignore-every LICENSELINT
+/**
+ * Copyright (c) Meta Platforms, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * Metal IVF-PQ index: 8-bit product quantization with precomputed
+ * per-query lookup tables and GPU IVF list scanning.
+ */
+
+#pragma once
+
+#import <Metal/Metal.h>
+
+#include <faiss/IndexIVFPQ.h>
+#include <faiss/gpu_metal/MetalIndex.h>
+
+#include <memory>
+
+namespace faiss {
+namespace gpu_metal {
+
+class MetalIVFPQImpl;
+
+class MetalIndexIVFPQ : public MetalIndex {
+   public:
+    MetalIndexIVFPQ(
+            std::shared_ptr<MetalResources> resources,
+            int dims,
+            idx_t nlist,
+            int M,
+            int nbitsPerIdx,
+            faiss::MetricType metric,
+            float metricArg = 0.0f,
+            MetalIndexConfig config = MetalIndexConfig());
+
+    MetalIndexIVFPQ(
+            std::shared_ptr<MetalResources> resources,
+            const faiss::IndexIVFPQ* cpuIndex,
+            MetalIndexConfig config = MetalIndexConfig());
+
+    ~MetalIndexIVFPQ() override;
+
+    void train(idx_t n, const float* x) override;
+    void add(idx_t n, const float* x) override;
+    void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
+    void reset() override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void copyFrom(const faiss::IndexIVFPQ* index);
+    void copyTo(faiss::IndexIVFPQ* index) const;
+
+    void updateQuantizer();
+    std::vector<idx_t> getListIndices(idx_t listId) const;
+    void reclaimMemory();
+
+    /// Pre-allocate GPU storage for the given total number of vectors.
+    void reserveMemory(idx_t numVecs);
+
+    idx_t nlist() const;
+    size_t nprobe() const;
+    int getNumSubQuantizers() const;
+    void setUsePrecomputedTables(bool enable);
+    bool getUsePrecomputedTables() const;
+
+   private:
+    std::unique_ptr<faiss::IndexIVFPQ> cpuIndex_;
+    std::unique_ptr<MetalIVFPQImpl> gpuIvf_;
+
+    void verifyPQSettings_() const;
+
+    void encodeResidualAndAppend_(
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            const idx_t* xids);
+};
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/MetalKernels.h

@@ -19,13 +19,15 @@
 namespace faiss {
 namespace gpu_metal {
 
+enum class IVFScanVariant { Standard, Small, Interleaved };
+
 class MetalKernels {
    public:
     explicit MetalKernels(id<MTLDevice> device);
     ~MetalKernels();
 
     bool isValid() const;
-    static constexpr int kMaxK = 256;
+    static constexpr int kMaxK = 2048;
 
     void encodeDistanceMatrix(
             id<MTLComputeCommandEncoder> enc,
@@ -37,6 +39,23 @@ class MetalKernels {
             int d,
             MetricType metric);
 
+    void encodeL2WithNorms(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> queries,
+            id<MTLBuffer> vectors,
+            id<MTLBuffer> distances,
+            id<MTLBuffer> vecNorms,
+            int nq,
+            int nb,
+            int d);
+
+    void encodeComputeNorms(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> vectors,
+            id<MTLBuffer> norms,
+            int nb,
+            int d);
+
     void encodeTopKThreadgroup(
             id<MTLComputeCommandEncoder> enc,
             id<MTLBuffer> distances,
@@ -47,6 +66,31 @@ class MetalKernels {
             int k,
             bool wantMin);
 
+    void encodeIVFScanList(
+            id<MTLComputeCommandEncoder> enc,
+            IVFScanVariant variant,
+            id<MTLBuffer> queries,
+            id<MTLBuffer> codes,
+            id<MTLBuffer> ids,
+            id<MTLBuffer> listOffset,
+            id<MTLBuffer> listLength,
+            id<MTLBuffer> coarseAssign,
+            id<MTLBuffer> perListDist,
+            id<MTLBuffer> perListIdx,
+            id<MTLBuffer> paramsBuf,
+            int nq,
+            int nprobe,
+            id<MTLBuffer> ilCodesOffset = nil);
+
+    void encodeIVFMergeLists(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> perListDist,
+            id<MTLBuffer> perListIdx,
+            id<MTLBuffer> outDist,
+            id<MTLBuffer> outIdx,
+            id<MTLBuffer> paramsBuf,
+            int nq);
+
     static int selectTopKVariantIndex(int k);
 
    private:
@@ -56,8 +100,9 @@ class MetalKernels {
     id<MTLLibrary> library_;
     std::unordered_map<std::string, id<MTLComputePipelineState>> cache_;
 
-    static constexpr int kTopKVariantSizes[] = {32, 64, 128, 256};
-    static constexpr int kNumTopKVariants = 4;
+    static constexpr int kTopKVariantSizes[] =
+            {32, 64, 128, 256, 512, 1024, 2048};
+    static constexpr int kNumTopKVariants = 7;
 };
 
 MetalKernels& getMetalKernels(id<MTLDevice> device);