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

Functions

nk_maxsim_pack_size_*_best× 3Returns packed buffer size in bytes for a maxsim vector set.

maxsim.h:81 ↗

nk_status_t nk_maxsim_pack_size_bf16_best(
    nk_size_t vector_count,
    nk_size_t depth,
    nk_capability_t capabilities,
    nk_size_t *bytes);

Symbol	Backends
`nk_maxsim_pack_size_bf16_best`	8
`nk_maxsim_pack_size_f32_best`	8
`nk_maxsim_pack_size_f16_best`	8

Backends

nk_maxsim_pack_size_bf16serialgenoasapphireamxhaswellalderv128relaxedneonsdotsme
nk_maxsim_pack_size_f32serialicelakesapphireamxhaswellalderv128relaxedneonsdotsme
nk_maxsim_pack_size_f16serialicelakesapphireamxhaswellalderv128relaxedneonsdotsme

Parameters

`vector_count`	in	The number of vectors to pack.
`depth`	in	The number of dimensions per vector.
`capabilities`	in	One device's capabilities, like `nk_cpu_capabilities_enabled` reports.
`bytes`	out	The size of the packed buffer.

Notes

The packed layout is backend-specific and must be produced by the matching pack function.

Statuses

nk_missing_kernel_k No capability in capabilities has it.

nk_maxsim_packed_shape_*_best× 3Reads a packed MaxSim buffer's shape from its header.

maxsim.h:101 ↗

nk_status_t nk_maxsim_packed_shape_bf16_best(
    void const *packed,
    nk_size_t *vectors,
    nk_size_t *depth,
    nk_capability_t capabilities,
    nk_stream_t stream);

Symbol	Backends
`nk_maxsim_packed_shape_bf16_best`	8
`nk_maxsim_packed_shape_f32_best`	8
`nk_maxsim_packed_shape_f16_best`	8

Backends

nk_maxsim_packed_shape_bf16serialgenoasapphireamxhaswellalderv128relaxedneonsdotsme
nk_maxsim_packed_shape_f32serialicelakesapphireamxhaswellalderv128relaxedneonsdotsme
nk_maxsim_packed_shape_f16serialicelakesapphireamxhaswellalderv128relaxedneonsdotsme

Parameters

`packed`	in	A buffer produced by the matching nk_maxsim_pack_bf16_best.
`vectors`	out	Receives the vector count.
`depth`	out	Receives the inner dimension.
`capabilities`	in	One device's capabilities, like `nk_cpu_capabilities_enabled` reports.
`stream`	in	Null on the CPU, or the GPU stream of that device to queue on.

Statuses

nk_missing_kernel_k No capability in capabilities has it.

nk_maxsim_pack_*_best× 3Packs vectors into a backend-specific layout for maxsim computation.

maxsim.h:124 ↗

nk_status_t nk_maxsim_pack_bf16_best(
    nk_bf16_t const *vectors,
    nk_size_t vector_count,
    nk_size_t depth,
    nk_size_t stride,
    void *packed,
    nk_capability_t capabilities,
    nk_stream_t stream);

Symbol	`vectors`	Backends
`nk_maxsim_pack_bf16_best`	`nk_bf16_t`	8
`nk_maxsim_pack_f32_best`	`nk_f32_t`	8
`nk_maxsim_pack_f16_best`	`nk_f16_t`	8

Backends

nk_maxsim_pack_bf16serialgenoasapphireamxhaswellalderv128relaxedneonsdotsme
nk_maxsim_pack_f32serialicelakesapphireamxhaswellalderv128relaxedneonsdotsme
nk_maxsim_pack_f16serialicelakesapphireamxhaswellalderv128relaxedneonsdotsme

Parameters

`vectors`	in	The input vectors in row-major order. Packing fuses no transposition, so depth-major vectors are transposed in a separate pass first.
`vector_count`	in	The number of vectors.
`depth`	in	The number of dimensions per vector.
`stride`	in	The row stride in bytes for the input vectors.
`packed`	out	The output packed buffer from nk_maxsim_pack_size_bf16_best.
`capabilities`	in	One device's capabilities, like `nk_cpu_capabilities_enabled` reports.
`stream`	in	Null on the CPU, or the GPU stream of that device to queue on.

Statuses

nk_missing_kernel_k No capability in capabilities has it.

nk_maxsim_packed_*_best× 3Computes angular distance late-interaction on pre-packed vectors.

maxsim.h:153 ↗

nk_status_t nk_maxsim_packed_bf16_best(
    void const *query_packed,
    void const *document_packed,
    nk_size_t query_count,
    nk_size_t document_count,
    nk_size_t depth,
    nk_f32_t *result,
    nk_capability_t capabilities,
    nk_stream_t stream);

Symbol	`result`	Backends
`nk_maxsim_packed_bf16_best`	`nk_f32_t`	8
`nk_maxsim_packed_f32_best`	`nk_f64_t`	8
`nk_maxsim_packed_f16_best`	`nk_f32_t`	8

Backends

nk_maxsim_packed_bf16serialgenoasapphireamxhaswellalderv128relaxedneonsdotsme
nk_maxsim_packed_f32serialicelakesapphireamxhaswellalderv128relaxedneonsdotsme
nk_maxsim_packed_f16serialicelakesapphireamxhaswellalderv128relaxedneonsdotsme

Returns Σᵢ minⱼ angular(qᵢ, dⱼ) where angular = 1 - dot / sqrt(||q||² × ||d||²), and 0 for two zero vectors. A NaN element in any query or document makes its angles, their minima and so the result NaN.

Parameters

`query_packed`	in	Packed query vectors (from nk_maxsim_pack_bf16_best).
`document_packed`	in	Packed document vectors (from nk_maxsim_pack_bf16_best).
`query_count`	in	Number of query vectors.
`document_count`	in	Number of document vectors.
`depth`	in	Number of dimensions per vector.
`result`	out	Pointer to store the sum of per-query minimum angular distances.
`capabilities`	in	One device's capabilities, like `nk_cpu_capabilities_enabled` reports.
`stream`	in	Null on the CPU, or the GPU stream of that device to queue on.

Statuses

nk_missing_kernel_k No capability in capabilities has it.

nk_maxsim_output_dtypeReturns the output dtype for MaxSim late-interaction.

maxsim.h:491 ↗

nk_dtype_t nk_maxsim_output_dtype(nk_dtype_t dtype);

nk_maxsim_error_boundReturns the error bound of MaxSim scores, per nk_accumulation_error_bound of their output.

maxsim.h:501 ↗

nk_f64_t nk_maxsim_error_bound(nk_dtype_t dtype);

nk_maxsim_find_kernelFinds the MaxSim kernel of kind for dtype, from the best of capabilities.

maxsim.h:517 ↗

nk_status_t nk_maxsim_find_kernel(
    nk_kernel_kind_t kind,
    nk_dtype_t dtype,
    nk_capability_t capabilities,
    nk_kernel_punned_t *kernel,
    nk_capability_t *capability);

Parameters

`kind`	in	The verb whose kernel to find.
`dtype`	in	The input type the kernel reads.
`capabilities`	in	The mask to pick the best capability from.
`kernel`	out	The kernel, or null when none of `capabilities` has it.
`capability`	out	The capability the kernel belongs to, or zero.

Statuses

nk_missing_kernel_k
nk_missing_library_k Header-only.