Skip to content

[multi vector] Add tiled MinMax MaxSim kernels #1394

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
juchen-ms (partychen) wants to merge 2 commits into from
Closed

[multi vector] Add tiled MinMax MaxSim kernels #1394

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a1db87c
Refactor tiled MinMax MaxSim and integrate benchmarks
partychen Oct 13, 2026
a79c823
Address MinMax MaxSim review feedback
partychen Oct 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion diskann-benchmark/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,19 @@ is running. This incremental saving allows benchmarks to be interrupted without
In addition to the machine-readable JSON output files, a (hopefully) helpful summary of the
results will be printed to `stdout`.

With `--features multi-vector`, select MinMax8 queries/MinMax4 documents using
`type: "multi-vector-op"` and `element_type: "minmax8"`;
see `example\multi-vector.json` and `perf_test_inputs\multi-vector-minmax8*.json`.

```powershell
cargo run --release --package diskann-benchmark --features multi-vector -- run --input-file diskann-benchmark\perf_test_inputs\multi-vector-minmax8.json --output-file minmax-results.json
```

Quantization and query preparation are outside the timed region. Each run repeatedly scores
the same document matrix and reports nanoseconds per query/document vector pair (`ns/IP`),
including per-call scratch work. These results use a different workload and statistic from
benchmarks that sweep many distinct documents and report microseconds per document.

### Streaming Runs
Running the benchmark on a streaming workload is similar to other registered benchmarks,
relying on the file formats and streaming runbooks of `big-ann-benchmarks`
Expand Down Expand Up @@ -504,4 +517,3 @@ for a successful or unsuccessful match with the enclosed value. Doing these two
error reporting in the event of a dispatch fail much easier for the user to understand and fix.

Refer to implementations within the benchmarking framework for what some of this may look like.

30 changes: 30 additions & 0 deletions diskann-benchmark/example/multi-vector.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,36 @@
{ "num_query_vectors": 64, "num_doc_vectors": 32, "dim": 264, "loops_per_measurement": 2, "num_measurements": 1 }
]
}
},
{
"type": "multi-vector-op",
"content": {
"element_type": "minmax8",
"isa": "auto",
"runs": [
{ "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 2, "num_measurements": 1 },
{ "num_query_vectors": 7, "num_doc_vectors": 13, "dim": 63, "loops_per_measurement": 2, "num_measurements": 1 },
{ "num_query_vectors": 8, "num_doc_vectors": 32, "dim": 64, "loops_per_measurement": 2, "num_measurements": 1 },
{ "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 65, "loops_per_measurement": 2, "num_measurements": 1 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 2, "num_measurements": 1 },
{ "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 2, "num_measurements": 1 }
]
}
},
{
"type": "multi-vector-op",
"content": {
"element_type": "minmax8",
"isa": "scalar",
"runs": [
{ "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 2, "num_measurements": 1 },
{ "num_query_vectors": 7, "num_doc_vectors": 13, "dim": 63, "loops_per_measurement": 2, "num_measurements": 1 },
{ "num_query_vectors": 8, "num_doc_vectors": 32, "dim": 64, "loops_per_measurement": 2, "num_measurements": 1 },
{ "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 65, "loops_per_measurement": 2, "num_measurements": 1 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 2, "num_measurements": 1 },
{ "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 2, "num_measurements": 1 }
]
}
}
]
}
47 changes: 47 additions & 0 deletions diskann-benchmark/perf_test_inputs/multi-vector-minmax8-neon.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
{
"search_directories": [],
"jobs": [
{
"type": "multi-vector-op",
"content": {
"element_type": "minmax8",
"isa": "neon",
"runs": [
{ "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 },
{ "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }
]
}
},
{
"type": "multi-vector-op",
"content": {
"element_type": "minmax8",
"isa": "scalar",
"runs": [
{ "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 },
{ "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }
]
}
}
]
}
47 changes: 47 additions & 0 deletions diskann-benchmark/perf_test_inputs/multi-vector-minmax8-v4.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
{
"search_directories": [],
"jobs": [
{
"type": "multi-vector-op",
"content": {
"element_type": "minmax8",
"isa": "x86-64-v4",
"runs": [
{ "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 },
{ "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }
]
}
},
{
"type": "multi-vector-op",
"content": {
"element_type": "minmax8",
"isa": "scalar",
"runs": [
{ "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 },
{ "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }
]
}
}
]
}
47 changes: 47 additions & 0 deletions diskann-benchmark/perf_test_inputs/multi-vector-minmax8.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
{
"search_directories": [],
"jobs": [
{
"type": "multi-vector-op",
"content": {
"element_type": "minmax8",
"isa": "auto",
"runs": [
{ "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 },
{ "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }
]
}
},
{
"type": "multi-vector-op",
"content": {
"element_type": "minmax8",
"isa": "scalar",
"runs": [
{ "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 },
{ "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 },
{ "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }
]
}
}
]
}
31 changes: 28 additions & 3 deletions diskann-benchmark/src/inputs/multi_vector.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

use std::num::NonZeroUsize;

use diskann_benchmark_runner::{utils::datatype::DataType, Checker, Input};
use diskann_benchmark_runner::{Checker, Input};
use diskann_quantization::multi_vector::MaxSimIsa;
use serde::{Deserialize, Serialize};

Expand Down Expand Up @@ -69,14 +69,36 @@ pub(crate) struct Run {
pub(crate) num_measurements: NonZeroUsize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum ElementType {
Float32,
Float16,
Int8,
/// MinMax8 queries and MinMax4 documents, quantized from f32 fixtures.
MinMax8,
}

impl std::fmt::Display for ElementType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(match self {
Self::Float32 => "float32",
Self::Float16 => "float16",
Self::Int8 => "int8",
Self::MinMax8 => "minmax8",
})
}
}

///////////////////////
// Multi-Vector Op //
///////////////////////

/// A complete multi-vector benchmark job.
#[derive(Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct MultiVectorOp {
pub(crate) element_type: DataType,
pub(crate) element_type: ElementType,
pub(crate) isa: BenchIsa,
pub(crate) runs: Vec<Run>,
}
Expand All @@ -95,6 +117,9 @@ impl Input for MultiVectorOp {
}

fn from_raw(raw: Self::Raw, _checker: &mut Checker) -> anyhow::Result<Self> {
if raw.element_type == ElementType::MinMax8 {
anyhow::ensure!(!raw.runs.is_empty(), "MinMax8 requires at least one run");
}
Ok(raw)
}

Expand Down Expand Up @@ -126,7 +151,7 @@ impl Input for MultiVectorOp {
];

Self {
element_type: DataType::Float32,
element_type: ElementType::Float32,
isa: BenchIsa::Auto,
runs,
}
Expand Down
Loading
Loading