MaxSim | NumKong for C++

MaxSim is the late-interaction primitive used by systems such as ColBERT. It is not generic matrix multiplication. It packs query and document token vectors into a scoring-specific layout and computes a late-interaction score.

auto [queries, queries_status] = nk::tensor<nk::bf16_t>::full({32, 128}, nk::bf16_t::one());
auto [docs, docs_status] = nk::tensor<nk::bf16_t>::full({192, 128}, nk::bf16_t::one());
if (nk::failed(queries_status) || nk::failed(docs_status)) return nk::failed(queries_status) ? queries_status : docs_status;

auto [q, q_status] = nk::packed_maxsim<nk::bf16_t>::make(queries.as_matrix_view());
auto [d, d_status] = nk::packed_maxsim<nk::bf16_t>::make(docs.as_matrix_view());
if (nk::failed(q_status) || nk::failed(d_status)) return nk::failed(q_status) ? q_status : d_status;
auto score = nk::maxsim(q, d);

packed_maxsim is allocator-aware in the same way as packed_matrix. Its footprint is exposed through size_bytes().

API

Types

packed_maxsimPre-packed vector set for MaxSim — ColBERT late-interaction.

matrix.hpp:541 ↗

template <numeric_dtype value_type_,
          typename allocator_type_ = aligned_allocator<char>>
class packed_maxsim

MaxSim computes Σᵢ minⱼ angular(qᵢ, dⱼ) using quantized i8 screening followed by full-precision refinement. Both queries and documents must be independently packed before calling maxsim().

Supported types: bf16_t, f32_t, f16_t.

Functions

maxsim_pack_sizeEstimates the memory requirements for a maxsim packed vector set.

matrix.hpp:173 ↗

template <numeric_dtype in_type_>
expected<std::size_t> maxsim_pack_size(
    std::size_t vector_count,
    std::size_t depth,
    nk_capability_t capabilities = default_capabilities());

Parameters

`vector_count`	in	Number of vectors to pack.
`depth`	in	Number of dimensions per vector.
`capabilities`	in	Capabilities to pick from, or zero for the C++ template.

Returns

Size in bytes for the packed buffer, or missing_kernel_k when no capability in capabilities packs in_type_.

maxsim_packPacks vectors into a backend-specific layout for maxsim computation.

matrix.hpp:203 ↗

template <numeric_dtype in_type_>
status_t maxsim_pack(
    typename in_type_::raw_t const *vectors,
    std::size_t vector_count,
    std::size_t depth,
    std::size_t stride,
    void *packed,
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr);

Parameters

`vectors`	in	Input vectors in row-major order. Packing fuses no transposition, so depth-major vectors are transposed in a separate pass first.
`vector_count`	in	Number of vectors.
`depth`	in	Number of dimensions per vector.
`stride`	in	Row stride in bytes for the input vectors.
`packed`	out	Output packed buffer from maxsim_pack_size.
`capabilities`	in	Capabilities to pick from, or zero for the C++ template.
`stream`	in	Null on the CPU, or the stream of the device `capabilities` describes.

maxsim_packedComputes angular distance late-interaction on pre-packed vectors.

maxsim.hpp:40 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::maxsim_result_t>
status_t maxsim_packed(
    void const *query_packed,
    void const *document_packed,
    std::size_t query_count,
    std::size_t document_count,
    std::size_t depth,
    result_type_ *result,
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr);

Returns Σᵢ minⱼ angular(qᵢ, dⱼ).

Parameters

`query_packed`	in	Packed query vectors.
`document_packed`	in	Packed document vectors.
`query_count`	in	Number of query vectors.
`document_count`	in	Number of document vectors.
`depth`	in	Number of dimensions per vector.
`result`	out	Sum of per-query minimum angular distances.
`capabilities`	in	Capabilities to pick from, or zero for the C++ template, which reads packs of the same zero mask.
`stream`	in	Null on the CPU, or the stream of the device `capabilities` describes.

maxsim_referenceExhaustive angular reference for testing: Σᵢ minⱼ angular(qᵢ, dⱼ).

maxsim.hpp:95 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::angular_result_t>
status_t maxsim_reference(
    typename in_type_::raw_t const *queries,
    std::size_t query_count,
    std::size_t query_stride,
    typename in_type_::raw_t const *documents,
    std::size_t document_count,
    std::size_t document_stride,
    std::size_t depth,
    result_type_ *result,
    nk_capability_t capabilities = default_capabilities());

Computes all pairwise angular distances and picks the minimum per query. Uses f64 accumulator for precision.

Parameters

`queries`	in	Query vectors in row-major order.
`query_count`	in	Number of query vectors.
`query_stride`	in	Row stride in bytes for query vectors.
`documents`	in	Document vectors in row-major order.
`document_count`	in	Number of document vectors.
`document_stride`	in	Row stride in bytes for document vectors.
`depth`	in	Number of dimensions per vector.
`result`	out	Pointer to store the sum of per-query minimum angular distances.
`capabilities`	in	Capabilities the distances pick from, or zero for the C++ template.

maxsimMaxSim: Σᵢ minⱼ angular(qᵢ, dⱼ) on pre-packed vectors.

maxsim.hpp:134 ↗

template <numeric_dtype value_type_>
value_type_::maxsim_result_t maxsim(
    packed_maxsim<value_type_> const & queries,
    packed_maxsim<value_type_> const & documents) noexcept;