TurboQuant native distance path and HNSW persistence support (WIP quality tuning)

config-model/src/main/java/com/yahoo/schema/document/Attribute.java

@@ -41,7 +41,7 @@
  */
 public final class Attribute implements Cloneable, Serializable {
 
-    public enum DistanceMetric { EUCLIDEAN, ANGULAR, GEODEGREES, INNERPRODUCT, HAMMING, PRENORMALIZED_ANGULAR, DOTPRODUCT }
+    public enum DistanceMetric { EUCLIDEAN, ANGULAR, GEODEGREES, INNERPRODUCT, HAMMING, PRENORMALIZED_ANGULAR, DOTPRODUCT, TURBOQUANT }
 
     // Remember to change hashCode and equals when you add new fields
 

config-model/src/test/java/com/yahoo/schema/AttributeSettingsTestCase.java

@@ -348,6 +348,7 @@ void distance_metric_is_propagated_to_attributes_config() throws ParseException
         // TODO Vespa 9: Remove 'innerproduct' as alias for 'prenormalized-angular'.
         assertDerivedDistanceMetric(AttributesConfig.Attribute.Distancemetric.INNERPRODUCT, "innerproduct");
         assertDerivedDistanceMetric(AttributesConfig.Attribute.Distancemetric.PRENORMALIZED_ANGULAR, "prenormalized-angular");
+        assertDerivedDistanceMetric(AttributesConfig.Attribute.Distancemetric.TURBOQUANT, "turboquant");
     }
 
     private void assertDerivedDistanceMetric(AttributesConfig.Attribute.Distancemetric.Enum expDistanceMetric,

configdefinitions/src/vespa/attributes.def

@@ -37,7 +37,7 @@ attribute[].maxuncommittedmemory long default=130000
 
 # The distance metric to use for nearest neighbor search.
 # Is only used when the attribute is a 1-dimensional indexed tensor.
-attribute[].distancemetric enum { EUCLIDEAN, ANGULAR, GEODEGREES, INNERPRODUCT, HAMMING, PRENORMALIZED_ANGULAR, DOTPRODUCT } default=EUCLIDEAN
+attribute[].distancemetric enum { EUCLIDEAN, ANGULAR, GEODEGREES, INNERPRODUCT, HAMMING, PRENORMALIZED_ANGULAR, DOTPRODUCT, TURBOQUANT } default=EUCLIDEAN
 
 # Configuration parameters for a hnsw index used together with a 1-dimensional indexed tensor for approximate nearest neighbor search.
 attribute[].index.hnsw.enabled bool default=false

integration/schema-language-server/language-server/src/main/java/ai/vespa/schemals/lsp/schema/completion/provider/FixedKeywordBodies.java

@@ -188,6 +188,8 @@ public CompletionItem getBodySnippet() {
         CompletionUtils.constructBasic("euclidean"),
         CompletionUtils.constructBasic("angular"),
         CompletionUtils.constructBasic("dotproduct"),
+        CompletionUtils.constructBasic("innerproduct"),
+        CompletionUtils.constructBasic("turboquant"),
         CompletionUtils.constructBasic("prenormalized-angular"),
         CompletionUtils.constructBasic("geodegrees"),
         CompletionUtils.constructBasic("hamming")

integration/tmgrammar/grammar/vespa-schema.tmLanguage.json

@@ -770,7 +770,7 @@
       ]
     },
     "enum-value-inline": {
-      "match": "(?<![a-zA-Z0-9_-])(match|rank|rank-type|sorting|bolding|stemming|summary-to|distance-metric|normalizing|function|locale|strength|order)(:\\s*)(token|word|exact|text|gram|prefix|substring|suffix|cased|uncased|literal|identity|tags|source|bolding|full|static|dynamic|tokens|matched-elements-only|angular|dotproduct|euclidean|prenormalized-angular|hamming|geodegrees|best|shortest|multiple|none|primary|secondary|tertiary|quaternary|identical|lowercase|raw|ascending|descending|always|on-demand|never|normal|contextual)(?![a-zA-Z0-9_-])",
+      "match": "(?<![a-zA-Z0-9_-])(match|rank|rank-type|sorting|bolding|stemming|summary-to|distance-metric|normalizing|function|locale|strength|order)(:\\s*)(token|word|exact|text|gram|prefix|substring|suffix|cased|uncased|literal|identity|tags|source|bolding|full|static|dynamic|tokens|matched-elements-only|angular|dotproduct|euclidean|innerproduct|prenormalized-angular|hamming|geodegrees|turboquant|best|shortest|multiple|none|primary|secondary|tertiary|quaternary|identical|lowercase|raw|ascending|descending|always|on-demand|never|normal|contextual)(?![a-zA-Z0-9_-])",
       "captures": {
         "1": {
           "name": "keyword.control.vespa"
@@ -818,7 +818,7 @@
         },
         {
           "name": "variable.other.enummember.vespa",
-          "match": "\\b(token|word|exact|text|gram|prefix|substring|suffix|cased|uncased|literal|identity|tags|source|bolding|full|static|dynamic|tokens|matched-elements-only|angular|dotproduct|euclidean|prenormalized-angular|hamming|geodegrees|best|shortest|multiple|none|primary|secondary|tertiary|quaternary|identical|lowercase|raw|ascending|descending|always|on-demand|never|normal|contextual)\\b"
+          "match": "\\b(token|word|exact|text|gram|prefix|substring|suffix|cased|uncased|literal|identity|tags|source|bolding|full|static|dynamic|tokens|matched-elements-only|angular|dotproduct|euclidean|innerproduct|prenormalized-angular|hamming|geodegrees|turboquant|best|shortest|multiple|none|primary|secondary|tertiary|quaternary|identical|lowercase|raw|ascending|descending|always|on-demand|never|normal|contextual)\\b"
         },
         {
           "include": "#schema-keywords"
@@ -878,7 +878,7 @@
         },
         {
           "name": "variable.other.enummember.vespa",
-          "match": "\\b(token|word|exact|text|gram|prefix|substring|suffix|cased|uncased|literal|identity|tags|source|bolding|full|static|dynamic|tokens|matched-elements-only|angular|dotproduct|euclidean|prenormalized-angular|hamming|geodegrees|best|shortest|multiple|none|primary|secondary|tertiary|quaternary|identical|lowercase|raw|ascending|descending|always|on-demand|never|normal|contextual)\\b"
+          "match": "\\b(token|word|exact|text|gram|prefix|substring|suffix|cased|uncased|literal|identity|tags|source|bolding|full|static|dynamic|tokens|matched-elements-only|angular|dotproduct|euclidean|innerproduct|prenormalized-angular|hamming|geodegrees|turboquant|best|shortest|multiple|none|primary|secondary|tertiary|quaternary|identical|lowercase|raw|ascending|descending|always|on-demand|never|normal|contextual)\\b"
         },
         {
           "include": "#schema-keywords"

searchlib/src/tests/attribute/attribute_header/attribute_header_test.cpp

@@ -79,6 +79,7 @@ TEST(AttributeHeaderTest, can_be_added_to_and_extracted_from_generic_header)
     verify_roundtrip_serialization(HnswIPO({16, 100, DistanceMetric::InnerProduct}));
     verify_roundtrip_serialization(HnswIPO({16, 100, DistanceMetric::PrenormalizedAngular}));
     verify_roundtrip_serialization(HnswIPO({16, 100, DistanceMetric::Hamming}));
+    verify_roundtrip_serialization(HnswIPO({16, 100, DistanceMetric::TurboQuant}));
     verify_roundtrip_serialization(HnswIPO());
 }
 

searchlib/src/tests/attribute/attributemanager/attributemanager_test.cpp

@@ -274,6 +274,7 @@ TEST(AttributeManagerTest, require_that_config_can_be_converted)
         expect_distance_metric(AttributesConfig::Attribute::Distancemetric::INNERPRODUCT, DistanceMetric::InnerProduct);
         expect_distance_metric(AttributesConfig::Attribute::Distancemetric::PRENORMALIZED_ANGULAR, DistanceMetric::PrenormalizedAngular);
         expect_distance_metric(AttributesConfig::Attribute::Distancemetric::DOTPRODUCT, DistanceMetric::Dotproduct);
+        expect_distance_metric(AttributesConfig::Attribute::Distancemetric::TURBOQUANT, DistanceMetric::TurboQuant);
     }
     { // hnsw index default params (enabled)
         CACA a;

searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp

@@ -408,6 +408,8 @@ const std::string test_dir = "test_data/";
 const std::string attr_name = test_dir + "my_attr";
 
 const std::string hnsw_max_squared_norm = "hnsw.max_squared_norm";
+const std::string hnsw_turbo_quant_version = "hnsw.turbo_quant.version";
+const std::string hnsw_turbo_quant_levels = "hnsw.turbo_quant.levels";
 
 struct FixtureTraits {
     bool use_dense_tensor_attribute = false;
@@ -416,6 +418,7 @@ struct FixtureTraits {
     bool use_mock_index = false;
     bool use_mmap_file_allocator = false;
     bool use_mips_distance = false;
+    bool use_turbo_quant_distance = false;
 
     FixtureTraits dense() && {
         use_dense_tensor_attribute = true;
@@ -457,6 +460,14 @@ struct FixtureTraits {
         return *this;
     }
 
+    FixtureTraits turbo_quant_hnsw() && {
+        use_dense_tensor_attribute = true;
+        enable_hnsw_index = true;
+        use_mock_index = false;
+        use_turbo_quant_distance = true;
+        return *this;
+    }
+
     FixtureTraits direct() && {
         use_dense_tensor_attribute = false;
         use_direct_tensor_attribute = true;
@@ -718,7 +729,12 @@ Fixture::Fixture(const std::string &typeSpec, FixtureTraits traits)
       _mmap_allocator_base_dir("mmap-file-allocator-factory-dir")
 {
     if (traits.enable_hnsw_index) {
-        auto dm = traits.use_mips_distance ? DistanceMetric::Dotproduct : DistanceMetric::Euclidean;
+        auto dm = DistanceMetric::Euclidean;
+        if (traits.use_turbo_quant_distance) {
+            dm = DistanceMetric::TurboQuant;
+        } else if (traits.use_mips_distance) {
+            dm = DistanceMetric::Dotproduct;
+        }
         _cfg.set_distance_metric(dm);
         _cfg.set_hnsw_index_params(HnswIndexParams(4, 20, dm));
     }
@@ -1459,6 +1475,11 @@ class DenseTensorAttributeMipsIndex : public Fixture {
     DenseTensorAttributeMipsIndex() : Fixture(vec_2d_spec, FixtureTraits().mips_hnsw()) {}
 };
 
+class DenseTensorAttributeTurboQuantIndex : public Fixture {
+public:
+    DenseTensorAttributeTurboQuantIndex() : Fixture(vec_2d_spec, FixtureTraits().turbo_quant_hnsw()) {}
+};
+
 TEST(TensorAttributeTest, Nearest_neighbor_index_with_mips_distance_metrics_stores_square_of_max_distance)
 {
     DenseTensorAttributeMipsIndex f;
@@ -1472,6 +1493,23 @@ TEST(TensorAttributeTest, Nearest_neighbor_index_with_mips_distance_metrics_stor
     EXPECT_EQ(130.0, norm_store.get_max());
 }
 
+TEST(TensorAttributeTest, Nearest_neighbor_index_with_turboquant_distance_metric_stores_square_of_max_distance)
+{
+    DenseTensorAttributeTurboQuantIndex f;
+    f.set_example_tensors();
+    EXPECT_TRUE(f.save());
+    auto header = f.get_file_header();
+    EXPECT_TRUE(header.hasTag(hnsw_max_squared_norm));
+    EXPECT_EQ(130.0, header.getTag(hnsw_max_squared_norm).asFloat());
+    EXPECT_TRUE(header.hasTag(hnsw_turbo_quant_version));
+    EXPECT_EQ(1, header.getTag(hnsw_turbo_quant_version).asInteger());
+    EXPECT_TRUE(header.hasTag(hnsw_turbo_quant_levels));
+    EXPECT_EQ(4, header.getTag(hnsw_turbo_quant_levels).asInteger());
+    EXPECT_TRUE(f.load());
+    auto& norm_store = dynamic_cast<MipsDistanceFunctionFactoryBase&>(f.hnsw_index().distance_function_factory()).get_max_squared_norm_store();
+    EXPECT_EQ(130.0, norm_store.get_max());
+}
+
 template <typename ParentT>
 class NearestNeighborBlueprintFixtureBase : public ParentT {
 private:

searchlib/src/tests/tensor/distance_functions/distance_functions_benchmark.cpp

@@ -5,6 +5,7 @@
 #include <vespa/searchlib/tensor/distance_functions.h>
 #include <vespa/searchlib/tensor/distance_function_factory.h>
 #include <vespa/searchlib/tensor/mips_distance_transform.h>
+#include <vespa/searchlib/tensor/turbo_quant_distance.h>
 #include <vespa/vespalib/util/benchmark_timer.h>
 #include <vespa/vespalib/util/classname.h>
 
@@ -86,6 +87,9 @@ void benchmark(size_t iterations, size_t elems, const std::string & dist_functio
     if (dist_functions.find("mips") != npos) {
         benchmark<T>(iterations, elems, MipsDistanceFunctionFactory<T>());
     }
+    if (dist_functions.find("turboquant") != npos) {
+        benchmark<T>(iterations, elems, TurboQuantDistanceFunctionFactory<T>());
+    }
 }
 
 void
@@ -108,7 +112,7 @@ int
 main(int argc, char *argv[]) {
     size_t num_iterations = 10000000;
     size_t num_elems = 1024;
-    std::string dist_functions = "angular euclid prenorm mips";
+    std::string dist_functions = "angular euclid prenorm mips turboquant";
     std::string data_types = "double float32 bfloat16 float8";
     if (argc > 1) { num_iterations = atol(argv[1]); }
     if (argc > 2) { num_elems = atol(argv[2]); }

searchlib/src/tests/tensor/distance_functions/distance_functions_test.cpp

@@ -5,8 +5,10 @@
 #include <vespa/searchlib/tensor/distance_functions.h>
 #include <vespa/searchlib/tensor/distance_function_factory.h>
 #include <vespa/searchlib/tensor/mips_distance_transform.h>
+#include <vespa/searchlib/tensor/turbo_quant_distance.h>
 #include <vespa/vespalib/gtest/gtest.h>
 #include <numbers>
+#include <random>
 #include <vector>
 
 #include <vespa/log/log.h>
@@ -764,6 +766,106 @@ TEST(DistanceFunctionsTest, dotproduct_can_reference_insertion_vector)
     expect_reference_insertion_vector<BFloat16>(0.0, DistanceMetric::Dotproduct, CellType::BFLOAT16);
 }
 
+TEST(DistanceFunctionsTest, turboquant_uses_dedicated_factory_type)
+{
+    auto dotproduct_factory = make_distance_function_factory(DistanceMetric::Dotproduct, CellType::FLOAT);
+    auto turboquant_factory = make_distance_function_factory(DistanceMetric::TurboQuant, CellType::FLOAT);
+    EXPECT_NE(nullptr, dynamic_cast<MipsDistanceFunctionFactory<float>*>(dotproduct_factory.get()));
+    EXPECT_NE(nullptr, dynamic_cast<TurboQuantDistanceFunctionFactory<float>*>(turboquant_factory.get()));
+    EXPECT_EQ(nullptr, dynamic_cast<TurboQuantDistanceFunctionFactory<float>*>(dotproduct_factory.get()));
+}
+
+TEST(DistanceFunctionsTest, turboquant_can_reference_insertion_vector)
+{
+    auto check_ref = [](CellType cell_type) {
+        std::vector<double> lhs{0.0, 1.0};
+        std::vector<double> rhs{0.0, 1.0};
+        auto factory = make_distance_function_factory(DistanceMetric::TurboQuant, cell_type);
+        auto func = factory->for_insertion_vector(t(lhs));
+        const double before = func->calc(t(rhs));
+        lhs[0] = 1.0;
+        lhs[1] = 0.0;
+        const double after = func->calc(t(rhs));
+        EXPECT_NE(before, after);
+    };
+    check_ref(CellType::FLOAT);
+    check_ref(CellType::DOUBLE);
+}
+
+TEST(DistanceFunctionsTest, turboquant_prefers_identical_vector_over_dissimilar_vector)
+{
+    std::vector<float> query{1.0f, 2.0f, 3.0f, 4.0f};
+    std::vector<float> same{1.0f, 2.0f, 3.0f, 4.0f};
+    std::vector<float> dissimilar{-4.0f, 3.0f, -2.0f, 1.0f};
+
+    auto factory = make_distance_function_factory(DistanceMetric::TurboQuant, CellType::FLOAT);
+    auto df = factory->for_query_vector(t(query));
+    const double same_distance = df->calc(t(same));
+    const double dissimilar_distance = df->calc(t(dissimilar));
+
+    EXPECT_LT(same_distance, dissimilar_distance);
+    EXPECT_GT(df->to_rawscore(same_distance), df->to_rawscore(dissimilar_distance));
+}
+
+TEST(DistanceFunctionsTest, turboquant_topk_overlap_with_exact_dotproduct_is_reasonable)
+{
+    constexpr size_t dims = 128;
+    constexpr size_t docs = 200;
+    constexpr size_t k = 10;
+
+    std::mt19937 rng(7);
+    std::uniform_real_distribution<float> dist(-3.0f, 3.0f);
+
+    std::vector<float> query(dims);
+    for (auto& value : query) {
+        value = dist(rng);
+    }
+    std::vector<std::vector<float>> vectors;
+    vectors.reserve(docs);
+    for (size_t i = 0; i < docs; ++i) {
+        std::vector<float> v(dims);
+        for (auto& value : v) {
+            value = dist(rng);
+        }
+        vectors.push_back(std::move(v));
+    }
+
+    auto turbo_factory = make_distance_function_factory(DistanceMetric::TurboQuant, CellType::FLOAT);
+    auto turbo_df = turbo_factory->for_query_vector(t(query));
+
+    auto exact_dot = [&query](const std::vector<float>& v) {
+        double sum = 0.0;
+        for (size_t i = 0; i < query.size(); ++i) {
+            sum += static_cast<double>(query[i]) * static_cast<double>(v[i]);
+        }
+        return sum;
+    };
+
+    std::vector<std::pair<double, uint32_t>> exact;
+    std::vector<std::pair<double, uint32_t>> approx;
+    exact.reserve(docs);
+    approx.reserve(docs);
+    for (uint32_t i = 0; i < docs; ++i) {
+        const auto& v = vectors[i];
+        exact.emplace_back(exact_dot(v), i);
+        approx.emplace_back(turbo_df->to_rawscore(turbo_df->calc(t(v))), i);
+    }
+    auto greater = [](const auto& lhs, const auto& rhs) { return lhs.first > rhs.first; };
+    std::partial_sort(exact.begin(), exact.begin() + k, exact.end(), greater);
+    std::partial_sort(approx.begin(), approx.begin() + k, approx.end(), greater);
+
+    size_t overlap = 0;
+    for (size_t i = 0; i < k; ++i) {
+        for (size_t j = 0; j < k; ++j) {
+            if (exact[i].second == approx[j].second) {
+                ++overlap;
+                break;
+            }
+        }
+    }
+    EXPECT_GE(overlap, 5u);
+}
+
 TEST(DistanceFunctionsTest, hamming_can_reference_insertion_vector)
 {
     expect_reference_insertion_vector<float>(2.0, DistanceMetric::Hamming, CellType::FLOAT);

searchlib/src/vespa/searchcommon/attribute/distance_metric.h

@@ -6,6 +6,6 @@
 
 namespace search::attribute {
 
-enum class DistanceMetric : uint8_t { Euclidean, Angular, GeoDegrees, InnerProduct, Hamming, PrenormalizedAngular, Dotproduct };
+enum class DistanceMetric : uint8_t { Euclidean, Angular, GeoDegrees, InnerProduct, Hamming, PrenormalizedAngular, Dotproduct, TurboQuant };
 
 }

searchlib/src/vespa/searchlib/attribute/configconverter.cpp

@@ -139,6 +139,9 @@ ConfigConverter::convert(const AttributesConfig::Attribute & cfg)
         case CfgDm::DOTPRODUCT:
             dm = DistanceMetric::Dotproduct;
             break;
+        case CfgDm::TURBOQUANT:
+            dm = DistanceMetric::TurboQuant;
+            break;
     }
     retval.set_distance_metric(dm);
     if (cfg.index.hnsw.enabled) {

searchlib/src/vespa/searchlib/attribute/distance_metric_utils.cpp

@@ -14,6 +14,7 @@ const std::string innerproduct = "innerproduct";
 const std::string prenormalized_angular = "prenormalized_angular";
 const std::string dotproduct = "dotproduct";
 const std::string hamming = "hamming";
+const std::string turboquant = "turboquant";
 
 }
 
@@ -28,6 +29,7 @@ DistanceMetricUtils::to_string(DistanceMetric metric)
         case DistanceMetric::Hamming: return hamming;
         case DistanceMetric::PrenormalizedAngular: return prenormalized_angular;
         case DistanceMetric::Dotproduct: return dotproduct;
+        case DistanceMetric::TurboQuant: return turboquant;
     }
     throw vespalib::IllegalArgumentException("Unknown distance metric " + std::to_string(static_cast<int>(metric)));
 }
@@ -49,6 +51,8 @@ DistanceMetricUtils::to_distance_metric(const std::string& metric)
         return DistanceMetric::Dotproduct;
     } else if (metric == hamming) {
         return DistanceMetric::Hamming;
+    } else if (metric == turboquant) {
+        return DistanceMetric::TurboQuant;
     } else {
         throw vespalib::IllegalStateException("Unknown distance metric '" + metric + "'");
     }

searchlib/src/vespa/searchlib/tensor/CMakeLists.txt

@@ -3,6 +3,7 @@ vespa_add_library(searchlib_tensor OBJECT
     SOURCES
     angular_distance.cpp
     mips_distance_transform.cpp
+    turbo_quant_distance.cpp
     bitvector_visited_tracker.cpp
     bound_distance_function.cpp
     default_nearest_neighbor_index_factory.cpp

searchlib/src/vespa/searchlib/tensor/distance_function_factory.cpp

@@ -3,6 +3,7 @@
 #include "distance_function_factory.h"
 #include "distance_functions.h"
 #include "mips_distance_transform.h"
+#include "turbo_quant_distance.h"
 
 using search::attribute::DistanceMetric;
 using vespalib::eval::CellType;
@@ -47,6 +48,14 @@ make_distance_function_factory(DistanceMetric variant, CellType cell_type)
                 case CellType::FLOAT:    return std::make_unique<MipsDistanceFunctionFactory<float>>(true);
                 default:                 return std::make_unique<MipsDistanceFunctionFactory<float>>();
             }
+        case DistanceMetric::TurboQuant:
+            switch (cell_type) {
+                case CellType::DOUBLE:   return std::make_unique<TurboQuantDistanceFunctionFactory<double>>(true);
+                case CellType::INT8:     return std::make_unique<TurboQuantDistanceFunctionFactory<Int8Float>>(true);
+                case CellType::BFLOAT16: return std::make_unique<TurboQuantDistanceFunctionFactory<vespalib::BFloat16>>(true);
+                case CellType::FLOAT:    return std::make_unique<TurboQuantDistanceFunctionFactory<float>>(true);
+                default:                 return std::make_unique<TurboQuantDistanceFunctionFactory<float>>();
+            }
         case DistanceMetric::GeoDegrees:
             return std::make_unique<GeoDistanceFunctionFactory>();
         case DistanceMetric::Hamming:

searchlib/src/vespa/searchlib/tensor/distance_function_factory.h

@@ -6,6 +6,8 @@
 #include "bound_distance_function.h"
 #include <vespa/searchcommon/attribute/distance_metric.h>
 
+namespace vespalib { class GenericHeader; }
+
 namespace search::tensor {
 
 /**
@@ -19,6 +21,8 @@ struct DistanceFunctionFactory {
     virtual ~DistanceFunctionFactory() = default;
     virtual BoundDistanceFunction::UP for_query_vector(TypedCells lhs) const = 0;
     virtual BoundDistanceFunction::UP for_insertion_vector(TypedCells lhs) const = 0;
+    virtual void save_state(vespalib::GenericHeader&) const {}
+    virtual void load_state(const vespalib::GenericHeader&) {}
     using UP = std::unique_ptr<DistanceFunctionFactory>;
 };