Matrix Kernels | NumKong for C++

This is a separate native subsystem from the raw vector kernels. It is the right tool when the right-hand side is reused many times.

#include <numkong/numkong.hpp>

namespace nk = ashvardanian::numkong;

auto [a, a_status] = nk::tensor<nk::f32_t>::full({2, 4}, nk::f32_t {1});
auto [b, b_status] = nk::tensor<nk::f32_t>::full({3, 4}, nk::f32_t {2});
if (nk::failed(a_status) || nk::failed(b_status)) return nk::failed(a_status) ? a_status : b_status;
auto [packed, packed_status] = nk::packed_matrix<nk::f32_t>::make(b.as_matrix_view());
if (nk::failed(packed_status)) return packed_status;

// Dot products, angular distances, and Euclidean distances all reuse the same packed B
auto [dots, dots_status] = nk::dots_packed<nk::f32_t>(a.as_matrix_view(), packed); // nk::matrix<nk::f64_t>
auto [angulars, angulars_status] = nk::angulars_packed<nk::f32_t>(a.as_matrix_view(), packed);
auto [euclideans, euclideans_status] = nk::euclideans_packed<nk::f32_t>(a.as_matrix_view(), packed);

This is GEMM-like in the workload shape, not in the strict BLAS API. The useful economics are:

one-time packing of B
one-time type preconversion where needed
depth padding handled internally
per-column norm reuse for angulars_packed and euclideans_packed
repeated reuse of the same packed RHS across many A batches

Caller-side alignment is not required. Owned packed_matrix storage uses its allocator. The C ABI also exposes nk_dots_pack_size_*_best, which writes the byte count of the capability the mask picks, so you can malloc the exact external buffer yourself. In C++, nk::dots_pack_size<nk::f32_t>(rows, depth) returns that count as an nk::expected<std::size_t>. A packed buffer records the capability that packed it, and a packed kernel of another capability refuses it with nk_pack_mismatch_k, so pack and multiply under the same mask.

§ Symmetric Kernels for SYRK-Like Workloads

The symmetric kernels solve a different problem. They compute self-similarity or self-distance without paying for both triangles independently.

auto [vectors, vectors_status] = nk::tensor<nk::f32_t>::full({100, 768}, nk::f32_t {1});
if (nk::failed(vectors_status)) return vectors_status;
auto [gram, gram_status] = nk::dots_symmetric<nk::f32_t>(vectors.as_matrix_view()); // nk::matrix<nk::f64_t>
auto [angular_dists, angular_status] = nk::angulars_symmetric<nk::f32_t>(vectors.as_matrix_view());
auto [euclidean_dists, euclidean_status] = nk::euclideans_symmetric<nk::f32_t>(vectors.as_matrix_view());

This is SYRK-like in the sense that the output is square and symmetric. The important difference from packed GEMM-style work is the partitioning model. You typically split by output row windows, not by distinct left batches against a shared packed right-hand side.

The arithmetic advantage is straightforward. The symmetric kernels avoid recomputing both (i, j) and (j, i) pairs. That cuts the pair count almost in half before any micro-kernel details matter.

§ Attention over Packed Key-Value Caches

Attention packs a ragged batch of keys and values once, then attends to it with any number of query batches. packed_attention keeps the segment offsets it was packed with, so the view overloads run self-attention over [tokens, heads, depth] queries without restating the batch, with query heads shared across fewer key-value heads.

#include <numkong/attention.hpp>

// Two sequences of 32 and 64 tokens, 8 query heads over 2 key-value heads of depth 64
auto [keys, keys_status] = nk::tensor<nk::bf16_t>::full({96, 2, 64}, nk::bf16_t::one());
auto [values, values_status] = nk::tensor<nk::bf16_t>::full({96, 2, 64}, nk::bf16_t::one());
auto [queries, queries_status] = nk::tensor<nk::bf16_t>::full({96, 8, 64}, nk::bf16_t::one());
auto [output, output_status] = nk::tensor<nk::f32_t>::zeros({96, 8, 64});
for (nk::status_t made : {keys_status, values_status, queries_status, output_status})
    if (nk::failed(made)) return made;
nk_u32_t const offsets[] = {0, 32, 96}, lengths[] = {32, 64};

auto [packed, packed_status] = nk::packed_attention<nk::bf16_t>::make(
    keys.view(), values.view(), nk::vector_view<nk_u32_t>(offsets, 3u), nk::vector_view<nk_u32_t>(lengths, 2u));
if (nk::failed(packed_status)) return packed_status;
std::size_t const all_keys = std::numeric_limits<std::size_t>::max(), causal = 0;
nk::status_t status = nk::attention_packed<nk::bf16_t>(queries.view(), packed, output.span(), 0.125f, all_keys, causal);
auto [fresh, fresh_status] = nk::attention_packed<nk::bf16_t>(queries.view(), packed, 0.125f); // nk::tensor<nk::f32_t>

Every overload takes the band of keys each query sees, keys_before and keys_after its position, every key by default. Queries align to the end of each segment’s keys, so keys_after = 0 is causal for prefill and decode alike, and keys_before = 4095 with it attends to at most 4096 keys per row. The raw-pointer overloads take the query offsets separately, for cross-attention, and a window over the query tokens × heads task grid, for sharding one launch across workers.

API

Classpacked_attentionPre-packed ragged KV-cache for scaled-dot-product attention.6

matrix.hpp:616 ↗

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

Packs K and V once, keeping a copy of the segment offsets after the packed bytes, so the view overloads of attention_packed() run self-attention against it without restating the batch.

Supported types: bf16_t, e4m3_t, i8_t.

packed_attention

matrix.hpp:627 ↗

packed_attention() noexcept=default;

packed_attention(allocator_type_ const & alloc) noexcept;

packed_attention(packed_attention && o) noexcept;

packed_attention(packed_attention const &)=delete;

~packed_*× 2

matrix.hpp:630 ↗

~packed_attention() noexcept;

Symbol
`~packed_attention`
`~packed_matrix`

operator=

matrix.hpp:639 ↗

packed_attention & operator=(packed_attention && o) noexcept;

packed_attention & operator=(packed_attention const &)=delete;

key_offsetsThe pack's slot boundaries, which the view overloads take as the query offsets.

matrix.hpp:724 ↗

vector_view<std::uint32_t> key_offsets() const noexcept;

key_lengthsThe keys each segment holds.

matrix.hpp:733 ↗

vector_view<std::uint32_t> key_lengths() const noexcept;

makeSize, allocate and pack a ragged batch of K and V tokens.

matrix.hpp:674 ↗

template <std::size_t max_rank_>
expected<packed_attention> make(
    tensor_view<value_type_, max_rank_> keys,
    tensor_view<value_type_, max_rank_> values,
    vector_view<std::uint32_t> key_offsets,
    vector_view<std::uint32_t> key_lengths = {},
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities()) noexcept;

Parameters

`keys`	in	[tokens,key_value_heads,depth] view whose heads sit back to back along its depth stride, so a transposed depth packs without a copy.
`values`	in	[tokens,key_value_heads,depth] view laid out like `keys`.
`key_offsets`	in	The slot boundaries of the segments, one more than there are segments, which must not decrease.
`key_lengths`	in	Keys each segment holds, at most its slot's width; zeros mark padding slots. Empty means every slot is full.
`alloc`	in	Allocator instance.
`capabilities`	in	Capabilities to pack for; attend with the same ones.

Returns

The pack.

Statuses

unexpected_dimensions_k The views or segments disagree or overrun the tokens.
missing_kernel_k
The packing kernel's failure.

Structpacked_matrixOwning, move-only, pre-packed matrix for efficient GEMM.8

matrix.hpp:443 ↗

template <numeric_dtype value_type_,
          typename allocator_type_ = aligned_allocator<char>>
struct packed_matrix

Wraps dots_pack to pre-arrange a matrix B into a cache-friendly layout. Use make() to create from a matrix_view, then pass to dots_packed() for computation.

packed_matrix

matrix.hpp:458 ↗

packed_matrix() noexcept=default;

packed_matrix(allocator_type_ const & alloc) noexcept;

packed_matrix(packed_matrix && other) noexcept;

packed_matrix(packed_matrix const &)=delete;

operator=

matrix.hpp:471 ↗

packed_matrix & operator=(packed_matrix && other) noexcept;

packed_matrix & operator=(packed_matrix const &)=delete;

rowsNumber of rows in the packed matrix (n).

matrix.hpp:517 ↗

size_type rows() const noexcept;

depthNumber of columns / depth (k).

matrix.hpp:520 ↗

size_type depth() const noexcept;

size_bytesSize of the packed buffer in bytes.

matrix.hpp:523 ↗

size_type size_bytes() const noexcept;

emptyTrue if no matrix is packed.

matrix.hpp:526 ↗

bool empty() const noexcept;

dataRaw pointer to the packed data.

matrix.hpp:529 ↗

void const * data() const noexcept;

makePack a 2D matrix_view into cache-efficient layout.

matrix.hpp:497 ↗

expected<packed_matrix> make(
    matrix_view<value_type_> b,
    allocator_type_ alloc = {}) noexcept;

Parameters

`b`	in	2D matrix view. Uses extents[0] as rows, extents[1] as depth.
`alloc`	in	Allocator instance.

Returns

The packed matrix.

Statuses

unexpected_dimensions_k A rank below 2.
missing_kernel_k Nothing packs value_type_.
The failure the packing kernel itself reports.

Functions

angulars_symmetric_workspace_sizeReturns the workspace bytes of angulars_symmetric, zero when no scratch is needed.

spatials.hpp:29 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::angular_result_t>
expected<std::size_t> angulars_symmetric_workspace_size(
    std::size_t vector_count,
    std::size_t depth,
    std::size_t row_start = 0,
    std::size_t row_count = std::numeric_limits<std::size_t>::max(),
    nk_capability_t capabilities = default_capabilities()) noexcept;

Parameters

`vector_count`	in	Number of vectors.
`depth`	in	Counts dimensions, a multiple of the values per byte.
`row_start`	in	First row of the result to compute, 0 by default.
`row_count`	in	Rows of the result to compute, all by default.
`capabilities`	in	Capabilities to pick from, or zero for the C++ template

euclideans_symmetric_workspace_sizeThe workspace bytes of euclideans_symmetric, queried like the angular one.

spatials.hpp:83 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::euclidean_result_t>
expected<std::size_t> euclideans_symmetric_workspace_size(
    std::size_t vector_count,
    std::size_t depth,
    std::size_t row_start = 0,
    std::size_t row_count = std::numeric_limits<std::size_t>::max(),
    nk_capability_t capabilities = default_capabilities()) noexcept;

angulars_packed_workspace_sizeThe workspace bytes of angulars_packed, queried like angulars_symmetric_workspace_size.

spatials.hpp:137 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::angular_result_t>
expected<std::size_t> angulars_packed_workspace_size(
    std::size_t row_count,
    std::size_t column_count,
    std::size_t depth,
    nk_capability_t capabilities = default_capabilities()) noexcept;

euclideans_packed_workspace_sizeThe workspace bytes of euclideans_packed, queried like the angular one.

spatials.hpp:175 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::euclidean_result_t>
expected<std::size_t> euclideans_packed_workspace_size(
    std::size_t row_count,
    std::size_t column_count,
    std::size_t depth,
    nk_capability_t capabilities = default_capabilities()) noexcept;

angulars_symmetricSymmetric angular distance matrix: C[i,j] = angular(A[i], A[j]).

spatials.hpp:238 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::angular_result_t>
status_t angulars_symmetric(
    in_type_ const *a,
    std::size_t vector_count,
    std::size_t depth,
    std::size_t a_stride,
    result_type_ *c,
    std::size_t c_stride,
    std::size_t row_start = 0,
    std::size_t row_count = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_, const_matrix_of<value_type_> input_matrix_, mutable_matrix_of<typename value_type_::angular_result_t> output_matrix_>
status_t angulars_symmetric(
    input_matrix_ const & input,
    output_matrix_ && output,
    std::size_t row_start = 0,
    std::size_t row_count = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Symmetric angular distances: C[i,j] = angular(A[i], A[j]), with the pointer overload's tail.

template <numeric_dtype value_type_, const_matrix_of<value_type_> input_matrix_, typename allocator_type_ = aligned_allocator<typename value_type_::angular_result_t>>
requires (!mutable_matrix_of<allocator_type_, typename value_type_::angular_result_t>)
expected<matrix<typename value_type_::angular_result_t, allocator_type_>> angulars_symmetric(
    input_matrix_ const & input,
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating symmetric angular distances, with any workspace in the result's block.

unexpected_dimensions_k unless the output is square over the input rows.

Parameters

`a`	in	Matrix A, shape [vector_count,depth]
`vector_count`	in	Number of vectors, n
`depth`	in	Counts dimensions, a multiple of the values per byte.
`a_stride`	in	Stride between vectors in A
`c`	out	Output matrix C, shape [n,n]
`c_stride`	in	Stride between rows of C in bytes
`row_start`	in	Starting row index, default 0
`row_count`	in	Number of rows to compute, default all
`workspace`	in	Scratch aligned to `nk_default_alignment_k`, sized by the matching query, or empty for the scratch-free path
`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

euclideans_symmetricSymmetric Euclidean distance matrix: C[i,j] = euclidean(A[i], A[j]).

spatials.hpp:344 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::euclidean_result_t>
status_t euclideans_symmetric(
    in_type_ const *a,
    std::size_t vector_count,
    std::size_t depth,
    std::size_t a_stride,
    result_type_ *c,
    std::size_t c_stride,
    std::size_t row_start = 0,
    std::size_t row_count = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_, const_matrix_of<value_type_> input_matrix_, mutable_matrix_of<typename value_type_::euclidean_result_t> output_matrix_>
status_t euclideans_symmetric(
    input_matrix_ const & input,
    output_matrix_ && output,
    std::size_t row_start = 0,
    std::size_t row_count = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Symmetric Euclidean distances: C[i,j] = euclidean(A[i], A[j]), trailing like the angular one.

template <numeric_dtype value_type_, const_matrix_of<value_type_> input_matrix_, typename allocator_type_ = aligned_allocator<typename value_type_::euclidean_result_t>>
requires (!mutable_matrix_of<allocator_type_, typename value_type_::euclidean_result_t>)
expected<matrix<typename value_type_::euclidean_result_t, allocator_type_>> euclideans_symmetric(
    input_matrix_ const & input,
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating symmetric Euclidean distances, with any workspace in the result's block.

unexpected_dimensions_k unless the output is square over the input rows.

Parameters

`a`	in	Matrix A, shape [vector_count,depth]
`vector_count`	in	Number of vectors, n
`depth`	in	Counts dimensions, a multiple of the values per byte.
`a_stride`	in	Stride between vectors in A
`c`	out	Output matrix C, shape [n,n]
`c_stride`	in	Stride between rows of C in bytes
`row_start`	in	Starting row index, default 0
`row_count`	in	Number of rows to compute, default all
`workspace`	in	Scratch aligned to `nk_default_alignment_k`, sized by the matching query, or empty for the scratch-free path
`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

angulars_packedPacked angular distances: C = angular(A, B_packed).

spatials.hpp:444 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::angular_result_t>
status_t angulars_packed(
    in_type_ const *a,
    void const *b_packed,
    result_type_ *c,
    std::size_t row_count,
    std::size_t column_count,
    std::size_t depth,
    std::size_t a_stride,
    std::size_t c_stride,
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_, packed_matrix_like packed_type_, const_matrix_of<value_type_> input_matrix_, mutable_matrix_of<typename value_type_::angular_result_t> output_matrix_>
status_t angulars_packed(
    input_matrix_ const & a,
    packed_type_ const & packed_b,
    output_matrix_ && c,
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Packed angular distances: C = angular(A, B_packed), taking the pointer overload's trailing parameters.

template <numeric_dtype value_type_, packed_matrix_like packed_type_, const_matrix_of<value_type_> input_matrix_, typename allocator_type_ = aligned_allocator<typename value_type_::angular_result_t>>
requires (!mutable_matrix_of<allocator_type_, typename value_type_::angular_result_t>)
expected<matrix<typename value_type_::angular_result_t, allocator_type_>> angulars_packed(
    input_matrix_ const & a,
    packed_type_ const & packed_b,
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating packed angular distances into one block with any workspace.

unexpected_dimensions_k for an empty pack, a rank below 2 or a shape mismatch.

Parameters

`a`	in	Matrix A, shape [row_count,depth]
`b_packed`	in	Packed B matrix, produced by nk_dots_pack_*
`c`	out	Output matrix C, shape [row_count,column_count]
`row_count`	in	Rows of A and C, m
`column_count`	in	Columns of B and C, n
`depth`	in	Shared inner dimension, k, in multiples of the values per byte.
`a_stride`	in	Stride between rows of A in bytes
`c_stride`	in	Stride between rows of C in bytes
`workspace`	in	Scratch aligned to `nk_default_alignment_k`, sized by the matching query, or empty for the scratch-free path
`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

euclideans_packedPacked Euclidean distances: C = euclidean(A, B_packed).

spatials.hpp:554 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::euclidean_result_t>
status_t euclideans_packed(
    in_type_ const *a,
    void const *b_packed,
    result_type_ *c,
    std::size_t row_count,
    std::size_t column_count,
    std::size_t depth,
    std::size_t a_stride,
    std::size_t c_stride,
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_, packed_matrix_like packed_type_, const_matrix_of<value_type_> input_matrix_, mutable_matrix_of<typename value_type_::euclidean_result_t> output_matrix_>
status_t euclideans_packed(
    input_matrix_ const & a,
    packed_type_ const & packed_b,
    output_matrix_ && c,
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Packed Euclidean distances: C = euclidean(A, B_packed), trailing like the angular one.

template <numeric_dtype value_type_, packed_matrix_like packed_type_, const_matrix_of<value_type_> input_matrix_, typename allocator_type_ = aligned_allocator<typename value_type_::euclidean_result_t>>
requires (!mutable_matrix_of<allocator_type_, typename value_type_::euclidean_result_t>)
expected<matrix<typename value_type_::euclidean_result_t, allocator_type_>> euclideans_packed(
    input_matrix_ const & a,
    packed_type_ const & packed_b,
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating packed Euclidean distances into one block with any workspace.

unexpected_dimensions_k for an empty pack, a rank below 2 or a mismatched shape.

Parameters

`a`	in	Matrix A, shape [row_count,depth]
`b_packed`	in	Packed B matrix, produced by nk_dots_pack_*
`c`	out	Output matrix C, shape [row_count,column_count]
`row_count`	in	Rows of A and C, m
`column_count`	in	Columns of B and C, n
`depth`	in	Shared inner dimension, k, in multiples of the values per byte.
`a_stride`	in	Stride between rows of A in bytes
`c_stride`	in	Stride between rows of C in bytes
`workspace`	in	Scratch aligned to `nk_default_alignment_k`, sized by the matching query, or empty for the scratch-free path
`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

dots_unpackedReference unpacked GEMM: C = A × Bᵀ (row-major A and B, B transposed).

dots.hpp:47 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::dot_result_t>
void dots_unpacked(
    operand_pointer<in_type_> a,
    operand_pointer<in_type_> b,
    result_type_ *c,
    std::size_t rows,
    std::size_t column_count,
    std::size_t depth,
    std::size_t a_stride,
    std::size_t b_stride,
    std::size_t c_stride) noexcept;

This matches BLAS sgemm/dgemm with CblasNoTrans for A and CblasTrans for B. Useful as a reference implementation for validating BLAS/MKL/Accelerate.

The scales of a block-scaled row sit at its code stride divided by the format's block bytes. Elements are multiplied by their block's scale and the operand's tensor scale, a missing one reading as 1, before they are multiplied and summed.

Parameters

`a`	in	Row-major matrix A of shape [rows,depth], see `operand_pointer`
`b`	in	Row-major matrix B of shape [columns,depth], accessed as Bᵀ
`c`	out	Row-major output matrix C of shape [rows,columns]
`rows`	in	Rows of A and C (m)
`column_count`	in	Rows of B and columns of C (n)
`depth`	in	Columns of A and B (k). Counts dimensions, a multiple of the values per byte.
`a_stride`	in	Stride between rows of A in bytes
`b_stride`	in	Stride between rows of B in bytes
`c_stride`	in	Stride between rows of C in bytes

dots_unpacked_conjugatedConjugated unpacked dot products: C = A × Bᴴ (Hermitian inner product, row-major).

dots.hpp:117 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::dot_result_t>
void dots_unpacked_conjugated(
    in_type_ const *a,
    in_type_ const *b,
    result_type_ *c,
    std::size_t rows,
    std::size_t column_count,
    std::size_t depth,
    std::size_t a_stride,
    std::size_t b_stride,
    std::size_t c_stride) noexcept;

Same as dots_unpacked, but conjugates elements of B before multiplication. For real types this is identical to dots_unpacked. For complex types this computes the standard Hermitian inner product matching cblas_{c,z}gemm with CblasConjTrans.

dots_packed_workspace_sizeReturns the workspace bytes of dots_packed, zero when no scratch is needed.

dots.hpp:145 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::dot_result_t>
expected<std::size_t> dots_packed_workspace_size(
    std::size_t rows,
    std::size_t column_count,
    std::size_t depth,
    nk_capability_t capabilities = default_capabilities()) noexcept;

Parameters

`rows`	in	Rows of A and C
`column_count`	in	Columns of B and C
`depth`	in	Columns of A, rows of B
`capabilities`	in	Capabilities to pick from, or zero for the C++ template

dots_symmetric_workspace_sizeReturns the workspace bytes of dots_symmetric, zero when no scratch is needed.

dots.hpp:190 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::dot_result_t>
expected<std::size_t> dots_symmetric_workspace_size(
    std::size_t vector_count,
    std::size_t depth,
    std::size_t rows_begin = 0,
    std::size_t rows_end = std::numeric_limits<std::size_t>::max(),
    nk_capability_t capabilities = default_capabilities()) noexcept;

Parameters

`vector_count`	in	Number of vectors
`depth`	in	Counts dimensions
`rows_begin`	in	First row of C to compute
`rows_end`	in	Row of C after the last
`capabilities`	in	Capabilities to pick from, or zero for the C++ template

dots_packedPacked dot products (batch matrix multiply): C = A × B (row-major).

dots.hpp:264 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::dot_result_t>
status_t dots_packed(
    operand_pointer<in_type_> a,
    void const *b_packed,
    result_type_ *c,
    std::size_t rows,
    std::size_t column_count,
    std::size_t depth,
    std::size_t a_stride,
    std::size_t c_stride,
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_, packed_matrix_like packed_type_, const_matrix_of<value_type_> input_matrix_, mutable_matrix_of<typename value_type_::dot_result_t> output_matrix_>
status_t dots_packed(
    input_matrix_ const & a,
    packed_type_ const & packed_b,
    output_matrix_ && c,
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Packed dot products: C = A × B_packedᵀ, taking the pointer overload's trailing parameters.

template <numeric_dtype value_type_, packed_matrix_like packed_type_, const_matrix_of<value_type_> input_matrix_, typename allocator_type_ = aligned_allocator<typename value_type_::dot_result_t>>
requires (!mutable_matrix_of<allocator_type_, typename value_type_::dot_result_t>)
expected<matrix<typename value_type_::dot_result_t, allocator_type_>> dots_packed(
    input_matrix_ const & a,
    packed_type_ const & packed_b,
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating packed dot products into one block with any workspace.

Block-scaled formats run only the C++ template, ignoring capabilities.

Parameters

`a`	in	Matrix A of shape [rows,depth]
`b_packed`	in	Packed matrix B of shape [depth,columns], with stride metadata appended
`c`	out	Output matrix C of shape [rows,columns]
`rows`	in	Rows of A and C
`column_count`	in	Columns of B and C
`depth`	in	Columns of A, rows of B, a multiple of the values per byte.
`a_stride`	in	Stride between rows of A in bytes
`c_stride`	in	Stride between rows of C in bytes
`workspace`	in	Scratch aligned to `nk_default_alignment_k`, sized by the matching query, or empty for the scratch-free path
`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

dots_symmetricSymmetric dot products: C = A × Aᵀ where C[i,j] = ⟨A[i], A[j]⟩

dots.hpp:361 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::dot_result_t>
status_t dots_symmetric(
    operand_pointer<in_type_> a,
    std::size_t vector_count,
    std::size_t depth,
    std::size_t a_stride,
    result_type_ *c,
    std::size_t c_stride,
    std::size_t rows_begin = 0,
    std::size_t rows_end = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_, const_matrix_of<value_type_> input_matrix_, mutable_matrix_of<typename value_type_::dot_result_t> output_matrix_>
status_t dots_symmetric(
    input_matrix_ const & input,
    output_matrix_ && output,
    std::size_t rows_begin = 0,
    std::size_t rows_end = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

C = A × Aᵀ where C[i,j] = ⟨A[i], A[j]⟩, taking the pointer overload's trailing parameters.

template <numeric_dtype value_type_, const_matrix_of<value_type_> input_matrix_, typename allocator_type_ = aligned_allocator<typename value_type_::dot_result_t>>
requires (!mutable_matrix_of<allocator_type_, typename value_type_::dot_result_t>)
expected<matrix<typename value_type_::dot_result_t, allocator_type_>> dots_symmetric(
    input_matrix_ const & input,
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating symmetric dot products: C = A × Aᵀ, with any workspace in the result's block.

Block-scaled formats run only the C++ template, ignoring capabilities.

Parameters

`a`	in	Matrix A of shape [vectors,depth]
`vector_count`	in	Number of vectors
`depth`	in	Counts dimensions, a multiple of the values per byte.
`a_stride`	in	Stride between vectors in A
`c`	out	Output matrix C of shape [vectors,vectors]
`c_stride`	in	Stride between rows of C in bytes
`rows_begin`	in	First row of C to compute, 0 by default
`rows_end`	in	Row of C after the last, clamped to `vector_count`, all by default
`workspace`	in	Scratch aligned to `nk_default_alignment_k`, sized by the matching query, or empty for the scratch-free path
`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

hammings_packed_workspace_sizeReturns the workspace bytes of hammings_packed, zero when no scratch is needed.

dots.hpp:444 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::hamming_result_t>
expected<std::size_t> hammings_packed_workspace_size(
    std::size_t rows,
    std::size_t column_count,
    std::size_t depth,
    nk_capability_t capabilities = default_capabilities()) noexcept;

hammings_symmetric_workspace_sizeReturns the workspace bytes of hammings_symmetric, zero when no scratch is needed.

dots.hpp:458 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::hamming_result_t>
expected<std::size_t> hammings_symmetric_workspace_size(
    std::size_t vector_count,
    std::size_t depth,
    std::size_t rows_begin = 0,
    std::size_t rows_end = std::numeric_limits<std::size_t>::max(),
    nk_capability_t capabilities = default_capabilities()) noexcept;

jaccards_packed_workspace_sizeReturns the workspace bytes of jaccards_packed, zero when no scratch is needed.

dots.hpp:476 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::jaccard_result_t>
expected<std::size_t> jaccards_packed_workspace_size(
    std::size_t rows,
    std::size_t column_count,
    std::size_t depth,
    nk_capability_t capabilities = default_capabilities()) noexcept;

jaccards_symmetric_workspace_sizeReturns the workspace bytes of jaccards_symmetric, zero when no scratch is needed.

dots.hpp:490 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::jaccard_result_t>
expected<std::size_t> jaccards_symmetric_workspace_size(
    std::size_t vector_count,
    std::size_t depth,
    std::size_t rows_begin = 0,
    std::size_t rows_end = std::numeric_limits<std::size_t>::max(),
    nk_capability_t capabilities = default_capabilities()) noexcept;

hammings_symmetricSymmetric Hamming distance matrix: C[i,j] = hamming(A[i], A[j]).

dots.hpp:528 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::hamming_result_t>
status_t hammings_symmetric(
    in_type_ const *a,
    std::size_t vector_count,
    std::size_t depth,
    std::size_t a_stride,
    result_type_ *c,
    std::size_t c_stride,
    std::size_t rows_begin = 0,
    std::size_t rows_end = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_, const_matrix_of<value_type_> input_matrix_, mutable_matrix_of<typename value_type_::hamming_result_t> output_matrix_>
status_t hammings_symmetric(
    input_matrix_ const & input,
    output_matrix_ && output,
    std::size_t rows_begin = 0,
    std::size_t rows_end = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Symmetric Hamming distances: C[i,j] = hamming(A[i], A[j]), trailing like dots_symmetric.

template <numeric_dtype value_type_, const_matrix_of<value_type_> input_matrix_, typename allocator_type_ = aligned_allocator<typename value_type_::hamming_result_t>>
requires (!mutable_matrix_of<allocator_type_, typename value_type_::hamming_result_t>)
expected<matrix<typename value_type_::hamming_result_t, allocator_type_>> hammings_symmetric(
    input_matrix_ const & input,
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating symmetric Hamming distances, with any workspace in the result's block.

Computes Hamming distances between all pairs of binary vectors. For u1x8_t inputs, distances are exact bit counts, returned as u32_t.

Parameters

`a`	in	Input matrix, shape [vector_count,depth]
`vector_count`	in	Number of vectors
`depth`	in	Counts dimensions, a multiple of the values per byte.
`a_stride`	in	Row stride in bytes
`c`	out	Output matrix, shape [vector_count,vector_count]
`c_stride`	in	Output row stride in bytes
`rows_begin`	in	First row to compute (default 0)
`rows_end`	in	Row after the last to compute, clamped to `vector_count` (default all)
`workspace`	in	At least what `hammings_symmetric_workspace_size` reports, or empty
`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

hammings_packedComputes Hamming distances between rows of A and columns of packed B.

dots.hpp:586 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::hamming_result_t>
status_t hammings_packed(
    in_type_ const *a,
    void const *b_packed,
    result_type_ *c,
    std::size_t rows,
    std::size_t column_count,
    std::size_t depth,
    std::size_t a_stride = 0,
    std::size_t c_stride = 0,
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_, packed_matrix_like packed_type_, const_matrix_of<value_type_> input_matrix_, mutable_matrix_of<typename value_type_::hamming_result_t> output_matrix_>
status_t hammings_packed(
    input_matrix_ const & a,
    packed_type_ const & packed_b,
    output_matrix_ && c,
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Packed Hamming distances: C = hamming(A, B_packed), trailing like dots_packed.

template <numeric_dtype value_type_, packed_matrix_like packed_type_, const_matrix_of<value_type_> input_matrix_, typename allocator_type_ = aligned_allocator<typename value_type_::hamming_result_t>>
requires (!mutable_matrix_of<allocator_type_, typename value_type_::hamming_result_t>)
expected<matrix<typename value_type_::hamming_result_t, allocator_type_>> hammings_packed(
    input_matrix_ const & a,
    packed_type_ const & packed_b,
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating packed Hamming distances into one block with any workspace.

Computes Hamming distances between binary vectors using optimized packed format. For u1x8_t inputs, distances are exact bit counts, returned as u32_t.

Parameters

`a`	in	Pointer to the first matrix (m x k).
`b_packed`	in	Pointer to the packed second matrix (k x n).
`c`	out	Pointer to the output matrix (m x n).
`rows`	in	Number of rows in A (m).
`column_count`	in	Number of columns in B (n).
`depth`	in	Counts dimensions, a multiple of the values per byte.
`a_stride`	in	Stride between consecutive rows of A in bytes.
`c_stride`	in	Stride between consecutive rows of C in bytes.
`workspace`	in	At least what `hammings_packed_workspace_size` reports, or empty.
`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.

jaccards_symmetricSymmetric Jaccard distance matrix: C[i,j] = jaccard(A[i], A[j]).

dots.hpp:638 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::jaccard_result_t>
status_t jaccards_symmetric(
    in_type_ const *a,
    std::size_t vector_count,
    std::size_t depth,
    std::size_t a_stride,
    result_type_ *c,
    std::size_t c_stride,
    std::size_t rows_begin = 0,
    std::size_t rows_end = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_, const_matrix_of<value_type_> input_matrix_, mutable_matrix_of<typename value_type_::jaccard_result_t> output_matrix_>
status_t jaccards_symmetric(
    input_matrix_ const & input,
    output_matrix_ && output,
    std::size_t rows_begin = 0,
    std::size_t rows_end = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Symmetric Jaccard distances: C[i,j] = jaccard(A[i], A[j]), trailing like dots_symmetric.

template <numeric_dtype value_type_, const_matrix_of<value_type_> input_matrix_, typename allocator_type_ = aligned_allocator<typename value_type_::jaccard_result_t>>
requires (!mutable_matrix_of<allocator_type_, typename value_type_::jaccard_result_t>)
expected<matrix<typename value_type_::jaccard_result_t, allocator_type_>> jaccards_symmetric(
    input_matrix_ const & input,
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating symmetric Jaccard distances, with any workspace in the result's block.

unexpected_dimensions_k unless the output is square over the input rows.

jaccards_packedComputes Jaccard distances between rows of A and columns of packed B.

dots.hpp:677 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::jaccard_result_t>
status_t jaccards_packed(
    in_type_ const *a,
    void const *b_packed,
    result_type_ *c,
    std::size_t rows,
    std::size_t column_count,
    std::size_t depth,
    std::size_t a_stride = 0,
    std::size_t c_stride = 0,
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_, packed_matrix_like packed_type_, const_matrix_of<value_type_> input_matrix_, mutable_matrix_of<typename value_type_::jaccard_result_t> output_matrix_>
status_t jaccards_packed(
    input_matrix_ const & a,
    packed_type_ const & packed_b,
    output_matrix_ && c,
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Packed Jaccard distances: C = jaccard(A, B_packed), trailing like dots_packed.

template <numeric_dtype value_type_, packed_matrix_like packed_type_, const_matrix_of<value_type_> input_matrix_, typename allocator_type_ = aligned_allocator<typename value_type_::jaccard_result_t>>
requires (!mutable_matrix_of<allocator_type_, typename value_type_::jaccard_result_t>)
expected<matrix<typename value_type_::jaccard_result_t, allocator_type_>> jaccards_packed(
    input_matrix_ const & a,
    packed_type_ const & packed_b,
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating packed Jaccard distances into one block with any workspace.

unexpected_dimensions_k for an empty pack, a rank below 2 or a mismatched shape.

dots_pack_sizeEstimates the memory requirements for packed B matrix.

matrix.hpp:46 ↗

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

Parameters

`columns`	in	Number of rows in B (n)
`depth`	in	Number of dimensions per row (k)
`capabilities`	in	Capabilities to pick from, or zero for the C++ template

Returns

Size in bytes for row-major B data plus stride metadata, or missing_kernel_k when no capability in capabilities packs in_type_

dots_packPacks matrix B into row-major form for efficient dots_packed access.

matrix.hpp:101 ↗

template <numeric_dtype in_type_>
status_t dots_pack(
    operand_pointer<in_type_> b,
    std::size_t columns,
    std::size_t depth,
    std::size_t b_stride,
    void *b_packed,
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr);

Block-scaled formats run only the C++ template, ignoring capabilities, and the pack holds the reference b points at, whose codes and scales must outlive it.

Parameters

`b`	in	Input matrix B, row-major, of shape [columns,depth], see `operand_pointer`; packing fuses no transposition, so a transposed B is transposed in a separate pass first
`columns`	in	Number of rows in B (n)
`depth`	in	Number of dimensions per row (k)
`b_stride`	in	Stride between rows of B in bytes
`b_packed`	out	Output buffer for packed row-major B with metadata
`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

attention_pack_sizeBounds the packed KV-cache of a ragged batch of segments, however its tokens split.

matrix.hpp:236 ↗

template <numeric_dtype in_type_>
expected<std::size_t> attention_pack_size(
    std::size_t key_value_head_count,
    std::size_t depth,
    std::size_t token_count,
    std::size_t segment_count,
    nk_capability_t capabilities = default_capabilities());

Parameters

`key_value_head_count`	in	Number of K/V heads, a nonzero divisor of the query head count.
`depth`	in	Head dimension; any value ≥ 1.
`token_count`	in	Live tokens across all segments.
`segment_count`	in	Number of segments packed together.
`capabilities`	in	Capabilities to pick from, or zero for the serial reference.

Returns

The size in bytes, or missing_kernel_k when no capability in capabilities packs in_type_.

attention_packPacks ragged K/V token matrices into a backend-opaque KV-cache blob.

matrix.hpp:325 ↗

template <numeric_dtype in_type_>
status_t attention_pack(
    std::size_t key_value_head_count,
    std::size_t depth,
    std::uint32_t const *key_offsets,
    std::uint32_t const *key_lengths,
    std::size_t segment_count,
    operand_pointer<in_type_> keys,
    std::size_t key_stride,
    operand_pointer<in_type_> values,
    std::size_t value_stride,
    void *key_value_packed,
    std::size_t tasks_begin = 0,
    std::size_t tasks_end = std::numeric_limits<std::size_t>::max(),
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr);

Parameters

`key_value_head_count`	in	K/V heads.
`depth`	in	Head dimension.
`key_offsets`	in	The `segment_count` + 1 slot boundaries, which must not decrease.
`key_lengths`	in	Keys each segment holds, at most its slot's width; zeros mark padding slots. Null means every slot is full.
`segment_count`	in	Number of segments in the batch.
`keys`	in	Token-major matrix, element (token, head, d) at `keys` + token × `key_stride` (head × depth + d) elements. Packing fuses no transposition, so depth-major K is transposed in a separate pass first.
`key_stride`	in	Bytes between tokens of `keys`.
`values`	in	Token-major matrix laid out like `keys` over `value_stride`. Depth-major V is likewise transposed in a separate pass first.
`value_stride`	in	Bytes between tokens of `values`.
`key_value_packed`	out	64-byte-aligned buffer of `attention_pack_size` bytes.
`tasks_begin`	in	First task of a window over the segments × K/V heads grid.
`tasks_end`	in	End of that half-open window, clipped to the grid, so callers can shard.
`capabilities`	in	Capabilities to pick from, or zero for the serial reference.
`stream`	in	Null on the CPU, or the stream of the device `capabilities` describes.

attention_packedRagged scaled-dot-product attention against a pre-packed KV-cache, under a band of keys.

attention.hpp:63 ↗

template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::attention_result_t>
status_t attention_packed(
    std::size_t head_count,
    std::size_t key_value_head_count,
    std::size_t depth,
    std::uint32_t const *query_offsets,
    std::size_t query_token_count,
    f32_t scale,
    std::size_t keys_before,
    std::size_t keys_after,
    operand_pointer<in_type_> queries,
    std::size_t queries_stride,
    void const *key_value_packed,
    result_type_ *output,
    std::size_t output_stride,
    f32_t *log_sum_exp,
    std::size_t tasks_begin = 0,
    std::size_t tasks_end = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr);

template <numeric_dtype value_type_, std::size_t max_rank_, typename allocator_type_>
status_t attention_packed(
    tensor_view<value_type_, max_rank_> queries,
    packed_attention<value_type_, allocator_type_> const & key_value_packed,
    tensor_span<typename value_type_::attention_result_t, max_rank_> output,
    f32_t scale,
    std::size_t keys_before = std::numeric_limits<std::size_t>::max(),
    std::size_t keys_after = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Self-attention of [tokens,heads,depth] queries against key_value_packed, whose pack-time key offsets split the query tokens too, under the keys_before, keys_after band, with the workspace, capabilities and stream of the pointer overload; unexpected_dimensions_k when the shapes disagree.

template <numeric_dtype value_type_, std::size_t max_rank_, typename packed_allocator_type_, typename allocator_type_ = aligned_allocator<typename value_type_::attention_result_t>>
expected<tensor<typename value_type_::attention_result_t, allocator_type_, max_rank_>> attention_packed(
    tensor_view<value_type_, max_rank_> queries,
    packed_attention<value_type_, packed_allocator_type_> const & key_value_packed,
    f32_t scale,
    std::size_t keys_before = std::numeric_limits<std::size_t>::max(),
    std::size_t keys_after = std::numeric_limits<std::size_t>::max(),
    allocator_type_ alloc = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

Allocating self-attention returning a fresh [tokens,heads,depth] tensor, with any workspace in its block; unexpected_dimensions_k when the shapes disagree, or else the failure of the allocation or of the kernel itself.

Queries align to the end of each segment's keys, as nk_attention_packed_bf16_best describes: causal sets keys_after to 0, and a sliding window of w keys sets keys_before to w − 1.

Parameters

`head_count`	in	Query heads.
`key_value_head_count`	in	K/V heads, at most `head_count`.
`depth`	in	Head dimension.
`query_offsets`	in	First query row of each segment, as segment count + 1 prefix sums.
`query_token_count`	in	Query rows the call addresses, `query_offsets[segments]`; GPU kernels size tensor maps and workspaces from it without reading the offsets.
`scale`	in	Score multiplier, typically 1 / √depth.
`keys_before`	in	Keys visible before each query's position, all of them for `max`.
`keys_after`	in	Keys visible after each query's position, all of them for `max`.
`queries`	in	Token-major matrix, one row of `head_count` × `depth` elements per token.
`queries_stride`	in	Row (token) stride of `queries` in bytes.
`key_value_packed`	in	A buffer `attention_pack` filled with the same `capabilities`.
`output`	out	Token-major matrix, one row of `head_count` × `depth` results per token.
`output_stride`	in	Row (token) stride of `output` in bytes.
`log_sum_exp`	out	Optional natural log-sum-exp per query token and head, for training.
`tasks_begin`	in	First task of a window over the query tokens × heads grid of output rows.
`tasks_end`	in	End of that half-open window, clipped to the grid. The window partitions the rows of `output` and `log_sum_exp` the call writes, per query token × head.
`workspace`	in	Scratch memory of `attention_packed_workspace_size` bytes on the device of `stream`, or empty for the path without scratch.
`capabilities`	in	Capabilities to pick from, or zero for the serial reference.
`stream`	in	Null on the CPU, or the stream of the device `capabilities` describes.

attention_packed_workspace_sizeReturns the workspace bytes attention_packed takes for its fast path on these shapes.

attention.hpp:193 ↗

template <numeric_dtype in_type_>
expected<std::size_t> attention_packed_workspace_size(
    std::size_t head_count,
    std::size_t key_value_head_count,
    std::size_t depth,
    std::size_t query_token_count,
    std::size_t keys_before,
    std::size_t keys_after,
    nk_capability_t capabilities = default_capabilities()) noexcept;

Parameters

`head_count`	in	Query heads.
`key_value_head_count`	in	K/V heads, at most `head_count`.
`depth`	in	Head dimension.
`query_token_count`	in	Query rows the call addresses.
`keys_before`	in	Keys visible before each query's position, all of them for `max`.
`keys_after`	in	Keys visible after each query's position, all of them for `max`.
`capabilities`	in	Capabilities to pick from, or zero for the serial reference.

attention_packed_gradientsGradients of attention_packed with respect to its queries, keys and values.

attention.hpp:275 ↗

template <numeric_dtype in_type_>
status_t attention_packed_gradients(
    std::size_t head_count,
    std::size_t key_value_head_count,
    std::size_t depth,
    std::uint32_t const *query_offsets,
    std::size_t query_token_count,
    f32_t scale,
    std::size_t keys_before,
    std::size_t keys_after,
    operand_pointer<in_type_> queries,
    std::size_t queries_stride,
    void const *key_value_packed,
    f32_t const *output,
    f32_t const *output_gradient,
    std::size_t output_stride,
    f32_t const *log_sum_exp,
    f32_t *query_gradient,
    std::size_t query_gradient_stride,
    f32_t *key_gradient,
    f32_t *value_gradient,
    std::size_t key_value_gradient_stride,
    std::size_t tasks_begin = 0,
    std::size_t tasks_end = std::numeric_limits<std::size_t>::max(),
    std::span<std::byte> workspace = {},
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr);

Every other parameter follows attention_packed; the band must be the forward's.

Parameters

`head_count`	in	Query heads.
`key_value_head_count`	in	K/V heads, at most `head_count`.
`depth`	in	Head dimension.
`query_offsets`	in	First query row of each segment, as in the forward pass.
`query_token_count`	in	Query rows the call addresses, as in the forward pass.
`scale`	in	Score multiplier of the forward pass.
`keys_before`	in	Keys before each query that it saw in the forward pass.
`keys_after`	in	Keys after each query that it saw in the forward pass.
`queries`	in	The forward's queries, token-major.
`queries_stride`	in	Row (token) stride of `queries` in bytes.
`key_value_packed`	in	The pack the forward pass read.
`output`	in	The forward's output for these queries and pack.
`output_gradient`	in	Gradient of the loss with respect to `output`, laid out like it.
`output_stride`	in	Row (token) stride of `output` and `output_gradient` in bytes.
`log_sum_exp`	in	The forward's log-sum-exp for these queries and pack.
`query_gradient`	out	Gradient with respect to `queries`, one row of `head_count` × `depth` per query token.
`query_gradient_stride`	in	Bytes between the `query_gradient` rows of consecutive tokens.
`key_gradient`	out	Gradient with respect to the packed keys, one row of `key_value_head_count` × `depth` per key token; rows past a segment's live keys are left unwritten by the kernel.
`value_gradient`	out	Gradient of the packed values, shaped like `key_gradient`.
`key_value_gradient_stride`	in	Bytes between the gradient rows of consecutive key tokens.
`tasks_begin`	in	First task of a window over the segments × key-value heads grid.
`tasks_end`	in	End of that half-open window, clipped to the grid. The window partitions the outputs the call writes: segments × K/V heads own their `key_gradient` and `value_gradient` rows and the `query_gradient` rows of their query heads.
`workspace`	in	Scratch memory of `attention_packed_gradients_workspace_size` bytes on the device of `stream`, or empty for the two-pass path.
`capabilities`	in	Capabilities to pick from, or zero for the serial reference.
`stream`	in	Null on the CPU, or the stream of the device `capabilities` describes.

attention_packed_gradients_workspace_sizeReturns the workspace bytes attention_packed_gradients takes for its single pass on these shapes, or 0 where it has none.

attention.hpp:399 ↗

template <numeric_dtype in_type_>
expected<std::size_t> attention_packed_gradients_workspace_size(
    std::size_t head_count,
    std::size_t key_value_head_count,
    std::size_t depth,
    std::size_t query_token_count,
    std::size_t keys_before,
    std::size_t keys_after,
    nk_capability_t capabilities = default_capabilities()) noexcept;

Parameters

`head_count`	in	Query heads.
`key_value_head_count`	in	K/V heads, at most `head_count`.
`depth`	in	Head dimension.
`query_token_count`	in	Query rows the call addresses.
`keys_before`	in	Keys before each query that it saw in the forward pass.
`keys_after`	in	Keys after each query that it saw in the forward pass.
`capabilities`	in	Capabilities to pick from, or zero for the serial reference.

attention_ropeNeoX split-half RoPE: rotates channel pairs of every head by per-token angles.

attention.hpp:450 ↗

template <numeric_dtype in_type_>
status_t attention_rope(
    in_type_ const *x,
    f32_t const *cos,
    f32_t const *sin,
    in_type_ *y,
    std::size_t rows,
    std::size_t head_count,
    std::size_t depth,
    std::size_t x_stride,
    std::size_t y_stride,
    nk_capability_t capabilities = default_capabilities(),
    nk_stream_t stream = nullptr) noexcept;

template <numeric_dtype value_type_>
status_t attention_rope(
    matrix_view<value_type_> x,
    vector_view<f32_t> cos,
    vector_view<f32_t> sin,
    matrix_span<value_type_> y,
    std::size_t head_count,
    std::size_t depth) noexcept;

NeoX split-half RoPE of a [rows,channels] matrix into a matching span, which may alias it, channels being head_count times an even depth; unexpected_dimensions_k when the shapes or the tables are too small.

Rotates each pair (i, i + depth / 2) of every head by the [rows, depth / 2] grids.

Parameters

`x`	in	Row-major matrix of shape [rows,channels]; channels = `head_count` · `depth`
`cos`	in	`[rows, depth / 2]` per-token cosine grid, shared across heads
`sin`	in	`[rows, depth / 2]` per-token sine grid, shared across heads
`y`	out	Output, same shape and dtype as x; may alias x for in-place rotation
`rows`	in	Token count
`head_count`	in	Heads per token
`depth`	in	Even channel count per head; channel `i` pairs with `i + depth / 2`
`x_stride`	in	Row (token) stride of x in bytes
`y_stride`	in	Row (token) stride of y in bytes
`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

TypesMember types5

`packed_matrix::value_type`	`value_type_`
`packed_matrix::result_type`	`typename value_type_::dot_result_t`
`packed_matrix::allocator_type`	`allocator_type_`
`packed_matrix::alloc_traits`	`std::allocator_traits<allocator_type_>`
`packed_matrix::size_type`	`std::size_t`