typedef nk_status_t(*) nk_dots_pack_size_punned_t(nk_size_t columns, nk_size_t depth, nk_size_t *bytes);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_packedandeuclideans_packed - repeated reuse of the same packed RHS across many
Abatches
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.
§ 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
Types
nk_dots_pack_size_punned_tPack sizes are host arithmetic, so they take no stream.
nk_sets_packed_punned_tBit vectors have no scaled form, so the set batches keep their unscaled shapes.
typedef nk_status_t(*) nk_sets_packed_punned_t(void const *a, void const *b_packed, void *c, nk_size_t rows, nk_size_t columns, nk_size_t depth, nk_size_t a_stride, nk_size_t c_stride, void *workspace, nk_size_t workspace_bytes, nk_stream_t stream);nk_attention_pack_size_punned_tPack sizes are host arithmetic, so they take no stream.
typedef nk_status_t(*) nk_attention_pack_size_punned_t(nk_size_t key_value_head_count, nk_size_t depth, nk_size_t token_count, nk_size_t segment_count, nk_size_t *bytes);nk_attention_packed_workspace_size_punned_tWorkspace sizes of attention and its gradients, host arithmetic like pack sizes.
typedef nk_status_t(*) nk_attention_packed_workspace_size_punned_t(nk_size_t head_count, nk_size_t key_value_head_count, nk_size_t depth, nk_size_t query_token_count, nk_size_t keys_before, nk_size_t keys_after, nk_size_t *bytes);nk_diagonal_band_tThe cells of a rows × columns grid a kernel visits, by their distance from the diagonal: column c of row r is visible when r − subdiagonals ≤ c ≤ r + superdiagonals, each side unbounded at NUMKONG_SIZE_MAX, like LAPACK's kl and ku.
struct nk_diagonal_band_tThe symmetric products visit the upper triangle (0, NUMKONG_SIZE_MAX), because A × Aᵀ mirrors it. Masked attention visits (keys_before, keys_after) around each query's position, because the other keys are forbidden rather than mirrored: causal attention, (NUMKONG_SIZE_MAX, 0), skips the same triangle of tiles for a different reason. So the band is a parameter of packed attention rather than a verb like symmetric, as its operands and outputs don't change.
Fields
subdiagonals | nk_size_t | |
|---|---|---|
superdiagonals | nk_size_t |
nk_diagonal_band_coverage_tWhere a banded grid's tile lies: no visible cell, every cell visible, or an edge across it.
enum nk_diagonal_band_coverage_tValues
nk_diagonal_band_outside_k | 0 | |
|---|---|---|
nk_diagonal_band_inside_k | 1 | |
nk_diagonal_band_crossing_k | 2 |
Functions
nk_dots_pack_size_*_best× 20Returns packed buffer size in bytes for second multiplier matrix (B).
nk_status_t nk_dots_pack_size_bf16_best(
nk_size_t columns,
nk_size_t depth,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_dots_pack_size_bf16_best | 24 |
nk_dots_pack_size_f16_best | 22 |
nk_dots_pack_size_e4m3_best | 22 |
nk_dots_pack_size_e5m2_best | 23 |
nk_dots_pack_size_e2m3_best | 24 |
nk_dots_pack_size_e2m1_best | 24 |
nk_dots_pack_size_nvfp4_best | 13 |
nk_dots_pack_size_mxfp4_best | 13 |
nk_dots_pack_size_mxfp6e2m3_best | 8 |
nk_dots_pack_size_mxfp6e3m2_best | 7 |
nk_dots_pack_size_mxfp8e4m3_best | 13 |
nk_dots_pack_size_mxfp8e5m2_best | 13 |
nk_dots_pack_size_e3m2_best | 20 |
nk_dots_pack_size_f32_best | 12 |
nk_dots_pack_size_f64_best | 12 |
nk_dots_pack_size_i8_best | 23 |
nk_dots_pack_size_u8_best | 23 |
nk_dots_pack_size_i4_best | 16 |
nk_dots_pack_size_u4_best | 16 |
nk_dots_pack_size_u1_best | 8 |
Backends
nk_dots_pack_size_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_size_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_size_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_size_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_size_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_size_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_size_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_pack_size_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_pack_size_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_dots_pack_size_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_dots_pack_size_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_pack_size_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_pack_size_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_size_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_pack_size_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_pack_size_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_pack_size_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_pack_size_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_pack_size_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_pack_size_u1serialsmebi32haswellicelakev128neonpowervsxloongsonasx
Parameters
columns | in | The number of rows in B (output columns). |
|---|---|---|
depth | in | Columns in B, counting dimensions, a multiple of the values per byte. |
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_kNo capability incapabilitieshas it.
nk_dots_packed_shape_*_best× 20Reads a packed B matrix's shape from its header.
nk_status_t nk_dots_packed_shape_bf16_best(
void const *b_packed,
nk_size_t *columns,
nk_size_t *depth,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | Backends |
|---|---|
nk_dots_packed_shape_bf16_best | 24 |
nk_dots_packed_shape_f16_best | 22 |
nk_dots_packed_shape_e4m3_best | 22 |
nk_dots_packed_shape_e5m2_best | 23 |
nk_dots_packed_shape_e2m3_best | 24 |
nk_dots_packed_shape_e2m1_best | 24 |
nk_dots_packed_shape_nvfp4_best | 13 |
nk_dots_packed_shape_mxfp4_best | 13 |
nk_dots_packed_shape_mxfp6e2m3_best | 8 |
nk_dots_packed_shape_mxfp6e3m2_best | 7 |
nk_dots_packed_shape_mxfp8e4m3_best | 13 |
nk_dots_packed_shape_mxfp8e5m2_best | 13 |
nk_dots_packed_shape_e3m2_best | 20 |
nk_dots_packed_shape_f32_best | 12 |
nk_dots_packed_shape_f64_best | 12 |
nk_dots_packed_shape_i8_best | 23 |
nk_dots_packed_shape_u8_best | 23 |
nk_dots_packed_shape_i4_best | 16 |
nk_dots_packed_shape_u4_best | 16 |
nk_dots_packed_shape_u1_best | 8 |
Backends
nk_dots_packed_shape_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_shape_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_shape_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_shape_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_shape_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_shape_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_shape_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_shape_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_shape_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_dots_packed_shape_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_dots_packed_shape_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_shape_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_shape_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_shape_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_packed_shape_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_packed_shape_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_shape_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_shape_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_shape_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_shape_u1serialsmebi32haswellicelakev128neonpowervsxloongsonasx
Parameters
b_packed | in | A buffer produced by the matching nk_dots_pack_bf16_best. |
|---|---|---|
columns | out | Receives the column count (B rows). |
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_kNo capability incapabilitieshas it.nk_device_code_mismatch_kA GPU stream capturing a graph, which reading the header back would void.
nk_dots_pack_*_best× 20Packs the second multiplier (B) matrix into a backend-specific layout.
nk_status_t nk_dots_pack_bf16_best(
nk_bf16_t const *b,
nk_size_t columns,
nk_size_t depth,
nk_size_t b_stride,
void *b_packed,
nk_size_t columns_begin,
nk_size_t columns_end,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | b | Backends |
|---|---|---|
nk_dots_pack_bf16_best | nk_bf16_t | 24 |
nk_dots_pack_f16_best | nk_f16_t | 22 |
nk_dots_pack_e4m3_best | nk_e4m3_t | 22 |
nk_dots_pack_e5m2_best | nk_e5m2_t | 23 |
nk_dots_pack_e2m3_best | nk_e2m3_t | 24 |
nk_dots_pack_e2m1_best | nk_e2m1x2_t | 24 |
nk_dots_pack_nvfp4_best | nk_nvfp4_cref_t | 13 |
nk_dots_pack_mxfp4_best | nk_mxfp4_cref_t | 13 |
nk_dots_pack_mxfp6e2m3_best | nk_mxfp6e2m3_cref_t | 8 |
nk_dots_pack_mxfp6e3m2_best | nk_mxfp6e3m2_cref_t | 7 |
nk_dots_pack_mxfp8e4m3_best | nk_mxfp8e4m3_cref_t | 13 |
nk_dots_pack_mxfp8e5m2_best | nk_mxfp8e5m2_cref_t | 13 |
nk_dots_pack_e3m2_best | nk_e3m2_t | 20 |
nk_dots_pack_f32_best | nk_f32_t | 12 |
nk_dots_pack_f64_best | nk_f64_t | 12 |
nk_dots_pack_i8_best | nk_i8_t | 23 |
nk_dots_pack_u8_best | nk_u8_t | 23 |
nk_dots_pack_i4_best | nk_i4x2_t | 16 |
nk_dots_pack_u4_best | nk_u4x2_t | 16 |
nk_dots_pack_u1_best | nk_u1x8_t | 8 |
Backends
nk_dots_pack_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_pack_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_pack_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_dots_pack_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_dots_pack_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_pack_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_pack_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_pack_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_pack_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_pack_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_pack_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_pack_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_pack_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_pack_u1serialsmebi32haswellicelakev128neonpowervsxloongsonasx
Parameters
b | in | The input B matrix in row-major order. Packing fuses no transposition, so a transposed B is transposed in a separate pass first. |
|---|---|---|
columns | in | The number of rows in B (output columns). |
depth | in | Columns in B, counting dimensions, a multiple of the values per byte. |
b_stride | in | The row stride in bytes for B. |
b_packed | out | The output packed buffer from nk_dots_pack_size_bf16_best. |
columns_begin | in | First output column to pack; 0 for a full pack. |
columns_end | in | One past the last output column to pack; columns for a full pack. |
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_kNo capability incapabilitieshas it.
Notes
Block-scaled B arrives as a reference to its first codes and scales. b_stride counts bytes of codes and whole blocks, and rows of scales sit b_stride / block_bytes apart, as nk_block_scaled_format_t says. The pack keeps B's tensor scale and every product applies it. depth must be a multiple of the block size: the CPU backends assert it, and the GPU backends return nk_unexpected_dimensions_k.
nk_dots_packed_workspace_size_*_best× 20Returns the scratch bytes that nk_dots_packed_bf16_best can use for these shapes.
nk_status_t nk_dots_packed_workspace_size_bf16_best(
nk_size_t rows,
nk_size_t columns,
nk_size_t depth,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_dots_packed_workspace_size_bf16_best | 24 |
nk_dots_packed_workspace_size_f16_best | 22 |
nk_dots_packed_workspace_size_e4m3_best | 22 |
nk_dots_packed_workspace_size_e5m2_best | 23 |
nk_dots_packed_workspace_size_e2m3_best | 24 |
nk_dots_packed_workspace_size_e2m1_best | 24 |
nk_dots_packed_workspace_size_nvfp4_best | 13 |
nk_dots_packed_workspace_size_mxfp4_best | 13 |
nk_dots_packed_workspace_size_mxfp6e2m3_best | 8 |
nk_dots_packed_workspace_size_mxfp6e3m2_best | 7 |
nk_dots_packed_workspace_size_mxfp8e4m3_best | 13 |
nk_dots_packed_workspace_size_mxfp8e5m2_best | 13 |
nk_dots_packed_workspace_size_e3m2_best | 20 |
nk_dots_packed_workspace_size_f32_best | 12 |
nk_dots_packed_workspace_size_f64_best | 12 |
nk_dots_packed_workspace_size_i8_best | 23 |
nk_dots_packed_workspace_size_u8_best | 23 |
nk_dots_packed_workspace_size_i4_best | 16 |
nk_dots_packed_workspace_size_u4_best | 16 |
nk_dots_packed_workspace_size_u1_best | 8 |
Backends
nk_dots_packed_workspace_size_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_workspace_size_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_workspace_size_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_workspace_size_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_workspace_size_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_workspace_size_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_workspace_size_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_workspace_size_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_workspace_size_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_dots_packed_workspace_size_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_dots_packed_workspace_size_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_workspace_size_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_workspace_size_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_workspace_size_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_packed_workspace_size_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_packed_workspace_size_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_workspace_size_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_workspace_size_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_workspace_size_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_workspace_size_u1serialsmebi32haswellicelakev128neonpowervsxloongsonasx
Parameters
rows | in | The number of rows in A. |
|---|---|---|
columns | in | The number of rows in B (output columns). |
depth | in | Columns in A and B, counting dimensions. |
capabilities | in | One device's capabilities, like nk_cpu_capabilities_enabled reports. |
bytes | out | The workspace size, zero when no scratch is needed. |
Statuses
nk_missing_kernel_kNo capability incapabilitieshas it.
nk_dots_packed_*_best× 20Computes C = A × Bᵀ using packed second multiplier matrix (B), accumulating into C.
nk_status_t nk_dots_packed_bf16_best(
nk_bf16_t const *a,
void const *b_packed,
nk_f32_t *c,
nk_size_t rows,
nk_size_t columns,
nk_size_t depth,
nk_size_t a_stride,
nk_size_t c_stride,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | a | c | Backends |
|---|---|---|---|
nk_dots_packed_bf16_best | nk_bf16_t | nk_f32_t | 24 |
nk_dots_packed_f16_best | nk_f16_t | nk_f32_t | 22 |
nk_dots_packed_e4m3_best | nk_e4m3_t | nk_f32_t | 22 |
nk_dots_packed_e5m2_best | nk_e5m2_t | nk_f32_t | 23 |
nk_dots_packed_e2m3_best | nk_e2m3_t | nk_f32_t | 24 |
nk_dots_packed_e2m1_best | nk_e2m1x2_t | nk_f32_t | 24 |
nk_dots_packed_nvfp4_best | nk_nvfp4_cref_t | nk_f32_t | 13 |
nk_dots_packed_mxfp4_best | nk_mxfp4_cref_t | nk_f32_t | 13 |
nk_dots_packed_mxfp6e2m3_best | nk_mxfp6e2m3_cref_t | nk_f32_t | 8 |
nk_dots_packed_mxfp6e3m2_best | nk_mxfp6e3m2_cref_t | nk_f32_t | 7 |
nk_dots_packed_mxfp8e4m3_best | nk_mxfp8e4m3_cref_t | nk_f32_t | 13 |
nk_dots_packed_mxfp8e5m2_best | nk_mxfp8e5m2_cref_t | nk_f32_t | 13 |
nk_dots_packed_e3m2_best | nk_e3m2_t | nk_f32_t | 20 |
nk_dots_packed_f32_best | nk_f32_t | nk_f64_t | 12 |
nk_dots_packed_f64_best | nk_f64_t | nk_f64_t | 12 |
nk_dots_packed_i8_best | nk_i8_t | nk_i32_t | 23 |
nk_dots_packed_u8_best | nk_u8_t | nk_u32_t | 23 |
nk_dots_packed_i4_best | nk_i4x2_t | nk_i32_t | 16 |
nk_dots_packed_u4_best | nk_u4x2_t | nk_u32_t | 16 |
nk_dots_packed_u1_best | nk_u1x8_t | nk_u32_t | 8 |
Backends
nk_dots_packed_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_dots_packed_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_dots_packed_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_packed_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_packed_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_packed_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_packed_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_packed_u1serialsmebi32haswellicelakev128neonpowervsxloongsonasx
Parameters
a | in | The input A matrix in row-major order. |
|---|---|---|
b_packed | in | The packed B matrix produced. |
c | out | The output C matrix in row-major order. |
rows | in | The number of rows in A. |
columns | in | The number of rows in B (output columns). |
depth | in | Columns in A and B, counting dimensions, a multiple of the values per byte. |
a_stride | in | The row stride in bytes for A. |
c_stride | in | The row stride in bytes for C. |
workspace | in | Scratch memory aligned to nk_default_alignment_k, device memory ordered by stream for GPU capabilities. Null with zero bytes selects the scratch-free path. |
workspace_bytes | in | At least what nk_dots_packed_workspace_size_bf16_best reports, or zero. Anything between zero and that size returns nk_unexpected_dimensions_k. |
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_misaligned_kA misalignedworkspace.nk_missing_kernel_kNo capability incapabilitieshas it.
Notes
Block-scaled A arrives as a reference, laid out as nk_dots_pack_nvfp4_best describes for B, and C holds true values, so every product takes the block scales and the tensor scales of both operands.
nk_dots_symmetric_workspace_size_*_best× 20Returns the scratch bytes that nk_dots_symmetric_bf16_best can use for these shapes.
nk_status_t nk_dots_symmetric_workspace_size_bf16_best(
nk_size_t vector_count,
nk_size_t depth,
nk_size_t rows_begin,
nk_size_t rows_end,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_dots_symmetric_workspace_size_bf16_best | 24 |
nk_dots_symmetric_workspace_size_f16_best | 22 |
nk_dots_symmetric_workspace_size_e4m3_best | 22 |
nk_dots_symmetric_workspace_size_e5m2_best | 23 |
nk_dots_symmetric_workspace_size_e2m3_best | 24 |
nk_dots_symmetric_workspace_size_e2m1_best | 24 |
nk_dots_symmetric_workspace_size_nvfp4_best | 13 |
nk_dots_symmetric_workspace_size_mxfp4_best | 13 |
nk_dots_symmetric_workspace_size_mxfp6e2m3_best | 8 |
nk_dots_symmetric_workspace_size_mxfp6e3m2_best | 7 |
nk_dots_symmetric_workspace_size_mxfp8e4m3_best | 13 |
nk_dots_symmetric_workspace_size_mxfp8e5m2_best | 13 |
nk_dots_symmetric_workspace_size_e3m2_best | 20 |
nk_dots_symmetric_workspace_size_f32_best | 12 |
nk_dots_symmetric_workspace_size_f64_best | 12 |
nk_dots_symmetric_workspace_size_i8_best | 23 |
nk_dots_symmetric_workspace_size_u8_best | 23 |
nk_dots_symmetric_workspace_size_i4_best | 16 |
nk_dots_symmetric_workspace_size_u4_best | 16 |
nk_dots_symmetric_workspace_size_u1_best | 8 |
Backends
nk_dots_symmetric_workspace_size_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_workspace_size_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_workspace_size_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_workspace_size_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_workspace_size_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_workspace_size_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_workspace_size_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_symmetric_workspace_size_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_symmetric_workspace_size_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_dots_symmetric_workspace_size_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_dots_symmetric_workspace_size_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_symmetric_workspace_size_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_symmetric_workspace_size_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_workspace_size_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_symmetric_workspace_size_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_symmetric_workspace_size_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_symmetric_workspace_size_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_symmetric_workspace_size_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_symmetric_workspace_size_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_symmetric_workspace_size_u1serialsmebi32haswellicelakev128neonpowervsxloongsonasx
Parameters
vector_count | in | Number of vectors (rows) in the input matrix. |
|---|---|---|
depth | in | Counts dimensions. |
rows_begin | in | First result row to compute. |
rows_end | in | Result row after the last to compute, clamped to vector_count. |
capabilities | in | One device's capabilities, like nk_cpu_capabilities_enabled reports. |
bytes | out | The workspace size, zero when no scratch is needed. |
Statuses
nk_missing_kernel_kNo capability incapabilitieshas it.
nk_dots_symmetric_*_best× 20Computes C = A × Aᵀ symmetric Gram matrix.
nk_status_t nk_dots_symmetric_bf16_best(
nk_bf16_t const *vectors,
nk_size_t vector_count,
nk_size_t depth,
nk_size_t stride,
nk_f32_t *result,
nk_size_t result_stride,
nk_size_t rows_begin,
nk_size_t rows_end,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | vectors | result | Backends |
|---|---|---|---|
nk_dots_symmetric_bf16_best | nk_bf16_t | nk_f32_t | 24 |
nk_dots_symmetric_f16_best | nk_f16_t | nk_f32_t | 22 |
nk_dots_symmetric_e4m3_best | nk_e4m3_t | nk_f32_t | 22 |
nk_dots_symmetric_e5m2_best | nk_e5m2_t | nk_f32_t | 23 |
nk_dots_symmetric_e2m3_best | nk_e2m3_t | nk_f32_t | 24 |
nk_dots_symmetric_e2m1_best | nk_e2m1x2_t | nk_f32_t | 24 |
nk_dots_symmetric_nvfp4_best | nk_nvfp4_cref_t | nk_f32_t | 13 |
nk_dots_symmetric_mxfp4_best | nk_mxfp4_cref_t | nk_f32_t | 13 |
nk_dots_symmetric_mxfp6e2m3_best | nk_mxfp6e2m3_cref_t | nk_f32_t | 8 |
nk_dots_symmetric_mxfp6e3m2_best | nk_mxfp6e3m2_cref_t | nk_f32_t | 7 |
nk_dots_symmetric_mxfp8e4m3_best | nk_mxfp8e4m3_cref_t | nk_f32_t | 13 |
nk_dots_symmetric_mxfp8e5m2_best | nk_mxfp8e5m2_cref_t | nk_f32_t | 13 |
nk_dots_symmetric_e3m2_best | nk_e3m2_t | nk_f32_t | 20 |
nk_dots_symmetric_f32_best | nk_f32_t | nk_f64_t | 12 |
nk_dots_symmetric_f64_best | nk_f64_t | nk_f64_t | 12 |
nk_dots_symmetric_i8_best | nk_i8_t | nk_i32_t | 23 |
nk_dots_symmetric_u8_best | nk_u8_t | nk_u32_t | 23 |
nk_dots_symmetric_i4_best | nk_i4x2_t | nk_i32_t | 16 |
nk_dots_symmetric_u4_best | nk_u4x2_t | nk_u32_t | 16 |
nk_dots_symmetric_u1_best | nk_u1x8_t | nk_u32_t | 8 |
Backends
nk_dots_symmetric_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_symmetric_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudaamperehopperblackwellmetalapple9apple10nk_dots_symmetric_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_dots_symmetric_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_dots_symmetric_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_symmetric_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudaamperehopperblackwellmetalapple9apple10nk_dots_symmetric_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_dots_symmetric_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_symmetric_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_dots_symmetric_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_symmetric_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_symmetric_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_symmetric_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_dots_symmetric_u1serialsmebi32haswellicelakev128neonpowervsxloongsonasx
Visits the upper triangle, the band (0, NUMKONG_SIZE_MAX) of nk_diagonal_band_t.
Parameters
vectors | in | Input matrix of row vectors in row-major order. |
|---|---|---|
vector_count | in | Number of vectors (rows) in the input matrix. |
depth | in | Counts dimensions, a multiple of the values per byte. |
stride | in | Row stride in bytes for the input matrix. |
result | out | Output symmetric matrix of vector_count × vector_count. |
result_stride | in | Row stride in bytes for the result matrix. |
rows_begin | in | First result row to compute, so threads split the work by rows. |
rows_end | in | Result row after the last to compute, clamped to vector_count. |
workspace | in | Scratch memory aligned to nk_default_alignment_k, device memory ordered by stream for GPU capabilities. Null with zero bytes selects the scratch-free path. |
workspace_bytes | in | At least what nk_dots_symmetric_workspace_size_bf16_best reports, or zero. Anything between zero and that size returns nk_unexpected_dimensions_k. |
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_misaligned_kA misalignedworkspace.nk_missing_kernel_kNo capability incapabilitieshas it.
nk_*_find_kernel× 4Finds the dots kernel of kind for dtype, from the best of capabilities.
nk_status_t nk_dots_find_kernel(
nk_kernel_kind_t kind,
nk_dtype_t dtype,
nk_capability_t capabilities,
nk_kernel_punned_t *kernel,
nk_capability_t *capability);| Symbol |
|---|
nk_dots_find_kernel |
nk_sets_find_kernel |
nk_spatials_find_kernel |
nk_attention_find_kernel |
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_knk_missing_library_kHeader-only.
nk_hammings_packed_u1_bestCompute Hamming distances between V rows and packed Q rows.
nk_status_t nk_hammings_packed_u1_best(
nk_u1x8_t const *v,
void const *query_packed,
nk_u32_t *result,
nk_size_t rows,
nk_size_t columns,
nk_size_t d,
nk_size_t v_stride,
nk_size_t r_stride,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);Backends
nk_hammings_packed_u1serialsmebi32haswellicelakeneonv128powervsxloongsonasx
Parameters
v | in | Input values matrix |
|---|---|---|
query_packed | in | Packed queries matrix |
result | out | Row-major results matrix |
rows | in | Number of rows in the results matrix |
columns | in | Number of columns in the results matrix |
d | in | Counts dimensions, a multiple of the values per byte. |
v_stride | in | Byte stride between rows of A |
r_stride | in | Byte stride between rows of C |
workspace | in | Scratch aligned to nk_default_alignment_k, or null with zero bytes. |
workspace_bytes | in | At least what the matching workspace-size query reports, or zero. |
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_kNo capability incapabilitieshas it.
nk_hammings_symmetric_u1_bestComputes C = A × Aᵀ symmetric Gram matrix of Hamming distances.
nk_status_t nk_hammings_symmetric_u1_best(
nk_u1x8_t const *vectors,
nk_size_t vector_count,
nk_size_t d,
nk_size_t stride,
nk_u32_t *result,
nk_size_t result_stride,
nk_size_t rows_begin,
nk_size_t rows_end,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);Backends
nk_hammings_symmetric_u1serialsmebi32haswellicelakeneonv128powervsxloongsonasx
Parameters
vectors | in | Input matrix of row vectors in row-major order. |
|---|---|---|
vector_count | in | Number of vectors (rows) in the input matrix. |
d | in | Counts dimensions, a multiple of the values per byte. |
stride | in | Row stride in bytes for the input matrix. |
result | out | Output symmetric matrix of vector_count × vector_count. |
result_stride | in | Row stride in bytes for the result matrix. |
rows_begin | in | First result row to compute, so threads split the work by rows. |
rows_end | in | Result row after the last to compute, clamped to vector_count. |
workspace | in | Scratch aligned to nk_default_alignment_k, or null with zero bytes. |
workspace_bytes | in | At least what the matching workspace-size query reports, or zero. |
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_kNo capability incapabilitieshas it.
nk_jaccards_packed_u1_bestCompute Jaccard distances between V rows and packed Q rows.
nk_status_t nk_jaccards_packed_u1_best(
nk_u1x8_t const *v,
void const *query_packed,
nk_f32_t *result,
nk_size_t rows,
nk_size_t columns,
nk_size_t d,
nk_size_t v_stride,
nk_size_t r_stride,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);Backends
nk_jaccards_packed_u1serialsmebi32haswellicelakeneonv128powervsxloongsonasx
Parameters
v | in | Input values matrix |
|---|---|---|
query_packed | in | Packed queries matrix (with norms) |
result | out | Row-major f32 results matrix |
rows | in | Number of rows in the results matrix |
columns | in | Number of columns in the results matrix |
d | in | Counts dimensions, a multiple of the values per byte. |
v_stride | in | Byte stride between rows of A |
r_stride | in | Byte stride between rows of C |
workspace | in | Scratch aligned to nk_default_alignment_k, or null with zero bytes. |
workspace_bytes | in | At least what the matching workspace-size query reports, or zero. |
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_kNo capability incapabilitieshas it.
nk_jaccards_symmetric_u1_bestComputes C = f(A, Aᵀ) symmetric Gram matrix of Jaccard distances.
nk_status_t nk_jaccards_symmetric_u1_best(
nk_u1x8_t const *vectors,
nk_size_t vector_count,
nk_size_t d,
nk_size_t stride,
nk_f32_t *result,
nk_size_t result_stride,
nk_size_t rows_begin,
nk_size_t rows_end,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);Backends
nk_jaccards_symmetric_u1serialsmebi32haswellicelakeneonv128powervsxloongsonasx
Parameters
vectors | in | Input matrix of row vectors in row-major order. |
|---|---|---|
vector_count | in | Number of vectors (rows). |
d | in | Counts dimensions, a multiple of the values per byte. |
stride | in | Row stride in bytes. |
result | out | Output symmetric f32 matrix of vector_count × vector_count. |
result_stride | in | Row stride in bytes for the result matrix. |
rows_begin | in | First result row to compute, so threads split the work by rows. |
rows_end | in | Result row after the last to compute, clamped to vector_count. |
workspace | in | Scratch aligned to nk_default_alignment_k, or null with zero bytes. |
workspace_bytes | in | At least what the matching workspace-size query reports, or zero. |
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_kNo capability incapabilitieshas it.
nk_*_packed_workspace_size_u1_best× 2Returns the workspace bytes of nk_hammings_packed_u1_best for these shapes.
nk_status_t nk_hammings_packed_workspace_size_u1_best(
nk_size_t rows,
nk_size_t columns,
nk_size_t d,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_hammings_packed_workspace_size_u1_best | 8 |
nk_jaccards_packed_workspace_size_u1_best | 8 |
Backends
nk_hammings_packed_workspace_size_u1serialsmebi32haswellicelakeneonv128powervsxloongsonasxnk_jaccards_packed_workspace_size_u1serialsmebi32haswellicelakeneonv128powervsxloongsonasx
Parameters
rows | in | Number of rows in the results matrix |
|---|---|---|
columns | in | Number of columns in the results matrix |
d | in | Counts dimensions, a multiple of the values per byte. |
capabilities | in | One device's capabilities, like nk_cpu_capabilities_enabled reports. |
bytes | out | The scratch the fast path takes, zero when the kernel needs none. |
nk_*_symmetric_workspace_size_u1_best× 2Returns the workspace bytes of nk_hammings_symmetric_u1_best for these shapes.
nk_status_t nk_hammings_symmetric_workspace_size_u1_best(
nk_size_t vector_count,
nk_size_t d,
nk_size_t rows_begin,
nk_size_t rows_end,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_hammings_symmetric_workspace_size_u1_best | 8 |
nk_jaccards_symmetric_workspace_size_u1_best | 8 |
Backends
nk_hammings_symmetric_workspace_size_u1serialsmebi32haswellicelakeneonv128powervsxloongsonasxnk_jaccards_symmetric_workspace_size_u1serialsmebi32haswellicelakeneonv128powervsxloongsonasx
Parameters
vector_count | in | Number of vectors (rows) in the input matrix. |
|---|---|---|
d | in | Counts dimensions, a multiple of the values per byte. |
rows_begin | in | First result row to compute. |
rows_end | in | Result row after the last to compute. |
capabilities | in | One device's capabilities, like nk_cpu_capabilities_enabled reports. |
bytes | out | The scratch the fast path takes, zero when the kernel needs none. |
nk_angulars_packed_workspace_size_*_best× 19Returns the workspace bytes of nk_angulars_packed_f32_best for one shape.
nk_status_t nk_angulars_packed_workspace_size_f32_best(
nk_size_t rows,
nk_size_t columns,
nk_size_t depth,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_angulars_packed_workspace_size_f32_best | 12 |
nk_angulars_packed_workspace_size_f64_best | 12 |
nk_angulars_packed_workspace_size_f16_best | 22 |
nk_angulars_packed_workspace_size_bf16_best | 24 |
nk_angulars_packed_workspace_size_e4m3_best | 22 |
nk_angulars_packed_workspace_size_e5m2_best | 23 |
nk_angulars_packed_workspace_size_e2m3_best | 24 |
nk_angulars_packed_workspace_size_e2m1_best | 24 |
nk_angulars_packed_workspace_size_nvfp4_best | 11 |
nk_angulars_packed_workspace_size_mxfp4_best | 11 |
nk_angulars_packed_workspace_size_mxfp6e2m3_best | 8 |
nk_angulars_packed_workspace_size_mxfp6e3m2_best | 7 |
nk_angulars_packed_workspace_size_mxfp8e4m3_best | 11 |
nk_angulars_packed_workspace_size_mxfp8e5m2_best | 11 |
nk_angulars_packed_workspace_size_e3m2_best | 20 |
nk_angulars_packed_workspace_size_i8_best | 23 |
nk_angulars_packed_workspace_size_u8_best | 23 |
nk_angulars_packed_workspace_size_i4_best | 16 |
nk_angulars_packed_workspace_size_u4_best | 16 |
Backends
nk_angulars_packed_workspace_size_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_angulars_packed_workspace_size_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_angulars_packed_workspace_size_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_workspace_size_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_workspace_size_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_workspace_size_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_workspace_size_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_workspace_size_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_workspace_size_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_angulars_packed_workspace_size_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_angulars_packed_workspace_size_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_angulars_packed_workspace_size_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_angulars_packed_workspace_size_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_angulars_packed_workspace_size_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_angulars_packed_workspace_size_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_workspace_size_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_packed_workspace_size_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_packed_workspace_size_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_packed_workspace_size_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10
Parameters
rows | in | Number of rows in A. |
|---|---|---|
columns | in | Number of columns in B (packed). |
depth | in | Shared inner dimension in dimensions. |
capabilities | in | One device's capabilities, like nk_cpu_capabilities_enabled reports. |
bytes | out | The scratch size, 0 when the capability needs none. |
Statuses
nk_missing_kernel_kNo capability incapabilitieshas it.
nk_angulars_symmetric_workspace_size_*_best× 19Returns the workspace bytes of nk_angulars_symmetric_f32_best for one shape.
nk_status_t nk_angulars_symmetric_workspace_size_f32_best(
nk_size_t vector_count,
nk_size_t depth,
nk_size_t rows_begin,
nk_size_t rows_end,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_angulars_symmetric_workspace_size_f32_best | 12 |
nk_angulars_symmetric_workspace_size_f64_best | 12 |
nk_angulars_symmetric_workspace_size_f16_best | 22 |
nk_angulars_symmetric_workspace_size_bf16_best | 24 |
nk_angulars_symmetric_workspace_size_e4m3_best | 22 |
nk_angulars_symmetric_workspace_size_e5m2_best | 23 |
nk_angulars_symmetric_workspace_size_e2m3_best | 24 |
nk_angulars_symmetric_workspace_size_e2m1_best | 24 |
nk_angulars_symmetric_workspace_size_nvfp4_best | 11 |
nk_angulars_symmetric_workspace_size_mxfp4_best | 11 |
nk_angulars_symmetric_workspace_size_mxfp6e2m3_best | 8 |
nk_angulars_symmetric_workspace_size_mxfp6e3m2_best | 7 |
nk_angulars_symmetric_workspace_size_mxfp8e4m3_best | 11 |
nk_angulars_symmetric_workspace_size_mxfp8e5m2_best | 11 |
nk_angulars_symmetric_workspace_size_e3m2_best | 20 |
nk_angulars_symmetric_workspace_size_i8_best | 23 |
nk_angulars_symmetric_workspace_size_u8_best | 23 |
nk_angulars_symmetric_workspace_size_i4_best | 16 |
nk_angulars_symmetric_workspace_size_u4_best | 16 |
Backends
nk_angulars_symmetric_workspace_size_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_angulars_symmetric_workspace_size_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_angulars_symmetric_workspace_size_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_workspace_size_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_workspace_size_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_workspace_size_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_workspace_size_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_workspace_size_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_workspace_size_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_angulars_symmetric_workspace_size_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_angulars_symmetric_workspace_size_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_angulars_symmetric_workspace_size_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_angulars_symmetric_workspace_size_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_angulars_symmetric_workspace_size_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_angulars_symmetric_workspace_size_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_workspace_size_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_symmetric_workspace_size_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_symmetric_workspace_size_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_symmetric_workspace_size_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10
Parameters
vector_count | in | Number of vectors (rows) in the input matrix. |
|---|---|---|
depth | in | Counts dimensions. |
rows_begin | in | First row of results to compute. |
rows_end | in | One past the last row of results to compute. |
capabilities | in | One device's capabilities, like nk_cpu_capabilities_enabled reports. |
bytes | out | The scratch size, 0 when the capability needs none. |
Statuses
nk_missing_kernel_kNo capability incapabilitieshas it.
nk_euclideans_packed_workspace_size_*_best× 19Returns the workspace bytes of nk_euclideans_packed_f32_best for one shape.
nk_status_t nk_euclideans_packed_workspace_size_f32_best(
nk_size_t rows,
nk_size_t columns,
nk_size_t depth,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_euclideans_packed_workspace_size_f32_best | 12 |
nk_euclideans_packed_workspace_size_f64_best | 12 |
nk_euclideans_packed_workspace_size_f16_best | 22 |
nk_euclideans_packed_workspace_size_bf16_best | 24 |
nk_euclideans_packed_workspace_size_e4m3_best | 22 |
nk_euclideans_packed_workspace_size_e5m2_best | 23 |
nk_euclideans_packed_workspace_size_e2m3_best | 24 |
nk_euclideans_packed_workspace_size_e2m1_best | 24 |
nk_euclideans_packed_workspace_size_nvfp4_best | 11 |
nk_euclideans_packed_workspace_size_mxfp4_best | 11 |
nk_euclideans_packed_workspace_size_mxfp6e2m3_best | 8 |
nk_euclideans_packed_workspace_size_mxfp6e3m2_best | 7 |
nk_euclideans_packed_workspace_size_mxfp8e4m3_best | 11 |
nk_euclideans_packed_workspace_size_mxfp8e5m2_best | 11 |
nk_euclideans_packed_workspace_size_e3m2_best | 20 |
nk_euclideans_packed_workspace_size_i8_best | 23 |
nk_euclideans_packed_workspace_size_u8_best | 23 |
nk_euclideans_packed_workspace_size_i4_best | 16 |
nk_euclideans_packed_workspace_size_u4_best | 16 |
Backends
nk_euclideans_packed_workspace_size_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_euclideans_packed_workspace_size_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_euclideans_packed_workspace_size_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_workspace_size_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_workspace_size_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_workspace_size_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_workspace_size_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_workspace_size_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_workspace_size_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_euclideans_packed_workspace_size_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_euclideans_packed_workspace_size_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_euclideans_packed_workspace_size_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_euclideans_packed_workspace_size_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_euclideans_packed_workspace_size_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_euclideans_packed_workspace_size_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_workspace_size_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_packed_workspace_size_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_packed_workspace_size_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_packed_workspace_size_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10
Parameters
rows | in | Number of rows in A. |
|---|---|---|
columns | in | Number of columns in B (packed). |
depth | in | Shared inner dimension in dimensions. |
capabilities | in | One device's capabilities, like nk_cpu_capabilities_enabled reports. |
bytes | out | The scratch size, 0 when the capability needs none. |
Statuses
nk_missing_kernel_kNo capability incapabilitieshas it.
nk_euclideans_symmetric_workspace_size_*_best× 19Returns the workspace bytes of nk_euclideans_symmetric_f32_best for one shape.
nk_status_t nk_euclideans_symmetric_workspace_size_f32_best(
nk_size_t vector_count,
nk_size_t depth,
nk_size_t rows_begin,
nk_size_t rows_end,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_euclideans_symmetric_workspace_size_f32_best | 12 |
nk_euclideans_symmetric_workspace_size_f64_best | 12 |
nk_euclideans_symmetric_workspace_size_f16_best | 22 |
nk_euclideans_symmetric_workspace_size_bf16_best | 24 |
nk_euclideans_symmetric_workspace_size_e4m3_best | 22 |
nk_euclideans_symmetric_workspace_size_e5m2_best | 23 |
nk_euclideans_symmetric_workspace_size_e2m3_best | 24 |
nk_euclideans_symmetric_workspace_size_e2m1_best | 24 |
nk_euclideans_symmetric_workspace_size_nvfp4_best | 11 |
nk_euclideans_symmetric_workspace_size_mxfp4_best | 11 |
nk_euclideans_symmetric_workspace_size_mxfp6e2m3_best | 8 |
nk_euclideans_symmetric_workspace_size_mxfp6e3m2_best | 7 |
nk_euclideans_symmetric_workspace_size_mxfp8e4m3_best | 11 |
nk_euclideans_symmetric_workspace_size_mxfp8e5m2_best | 11 |
nk_euclideans_symmetric_workspace_size_e3m2_best | 20 |
nk_euclideans_symmetric_workspace_size_i8_best | 23 |
nk_euclideans_symmetric_workspace_size_u8_best | 23 |
nk_euclideans_symmetric_workspace_size_i4_best | 16 |
nk_euclideans_symmetric_workspace_size_u4_best | 16 |
Backends
nk_euclideans_symmetric_workspace_size_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_euclideans_symmetric_workspace_size_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_euclideans_symmetric_workspace_size_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_workspace_size_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_workspace_size_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_workspace_size_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_workspace_size_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_workspace_size_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_workspace_size_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_euclideans_symmetric_workspace_size_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_euclideans_symmetric_workspace_size_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_euclideans_symmetric_workspace_size_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_euclideans_symmetric_workspace_size_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_euclideans_symmetric_workspace_size_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_euclideans_symmetric_workspace_size_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_workspace_size_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_symmetric_workspace_size_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_symmetric_workspace_size_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_symmetric_workspace_size_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10
Parameters
vector_count | in | Number of vectors (rows) in the input matrix. |
|---|---|---|
depth | in | Counts dimensions. |
rows_begin | in | First row of results to compute. |
rows_end | in | One past the last row of results to compute. |
capabilities | in | One device's capabilities, like nk_cpu_capabilities_enabled reports. |
bytes | out | The scratch size, 0 when the capability needs none. |
Statuses
nk_missing_kernel_kNo capability incapabilitieshas it.
nk_angulars_packed_*_best× 19Computes batched angular distances using a packed second matrix.
nk_status_t nk_angulars_packed_f32_best(
nk_f32_t const *a,
void const *b_packed,
nk_f64_t *result,
nk_size_t rows,
nk_size_t columns,
nk_size_t depth,
nk_size_t a_stride,
nk_size_t r_stride,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | a | result | Backends |
|---|---|---|---|
nk_angulars_packed_f32_best | nk_f32_t | nk_f64_t | 12 |
nk_angulars_packed_f64_best | nk_f64_t | nk_f64_t | 12 |
nk_angulars_packed_f16_best | nk_f16_t | nk_f32_t | 22 |
nk_angulars_packed_bf16_best | nk_bf16_t | nk_f32_t | 24 |
nk_angulars_packed_e4m3_best | nk_e4m3_t | nk_f32_t | 22 |
nk_angulars_packed_e5m2_best | nk_e5m2_t | nk_f32_t | 23 |
nk_angulars_packed_e2m3_best | nk_e2m3_t | nk_f32_t | 24 |
nk_angulars_packed_e2m1_best | nk_e2m1x2_t | nk_f32_t | 24 |
nk_angulars_packed_nvfp4_best | nk_nvfp4_cref_t | nk_f32_t | 11 |
nk_angulars_packed_mxfp4_best | nk_mxfp4_cref_t | nk_f32_t | 11 |
nk_angulars_packed_mxfp6e2m3_best | nk_mxfp6e2m3_cref_t | nk_f32_t | 8 |
nk_angulars_packed_mxfp6e3m2_best | nk_mxfp6e3m2_cref_t | nk_f32_t | 7 |
nk_angulars_packed_mxfp8e4m3_best | nk_mxfp8e4m3_cref_t | nk_f32_t | 11 |
nk_angulars_packed_mxfp8e5m2_best | nk_mxfp8e5m2_cref_t | nk_f32_t | 11 |
nk_angulars_packed_e3m2_best | nk_e3m2_t | nk_f32_t | 20 |
nk_angulars_packed_i8_best | nk_i8_t | nk_f32_t | 23 |
nk_angulars_packed_u8_best | nk_u8_t | nk_f32_t | 23 |
nk_angulars_packed_i4_best | nk_i4x2_t | nk_f32_t | 16 |
nk_angulars_packed_u4_best | nk_u4x2_t | nk_f32_t | 16 |
Backends
nk_angulars_packed_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_angulars_packed_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_angulars_packed_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_angulars_packed_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_angulars_packed_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_angulars_packed_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_angulars_packed_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_angulars_packed_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_angulars_packed_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_packed_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_packed_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_packed_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_packed_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10
Parameters
a | in | Input A matrix in row-major order. |
|---|---|---|
b_packed | in | Packed B with norms from nk_dots_pack_*_best on the same capabilities. |
result | out | Output matrix (rows x columns) of angular distances. |
rows | in | Number of rows in A. |
columns | in | Number of columns in B (packed). |
depth | in | Shared inner dimension in dimensions, a multiple of the values per byte. |
a_stride | in | Row stride in bytes for A. |
r_stride | in | Row stride in bytes for the result matrix. |
workspace | in | Scratch memory aligned to nk_default_alignment_k, in device memory ordered by stream on GPU capabilities. May be null with workspace_bytes of 0. |
workspace_bytes | in | At least what the matching workspace-size query reports for the fast path, or 0 for the scratch-free path. Anything in between is nk_unexpected_dimensions_k. |
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_misaligned_kA misalignedworkspace.nk_missing_kernel_kNo capability incapabilitieshas it.
nk_angulars_symmetric_*_best× 19Computes symmetric angular distance matrix (Gram-style) for a set of vectors.
nk_status_t nk_angulars_symmetric_f32_best(
nk_f32_t const *vectors,
nk_size_t vector_count,
nk_size_t depth,
nk_size_t stride,
nk_f64_t *result,
nk_size_t result_stride,
nk_size_t rows_begin,
nk_size_t rows_end,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | vectors | result | Backends |
|---|---|---|---|
nk_angulars_symmetric_f32_best | nk_f32_t | nk_f64_t | 12 |
nk_angulars_symmetric_f64_best | nk_f64_t | nk_f64_t | 12 |
nk_angulars_symmetric_f16_best | nk_f16_t | nk_f32_t | 22 |
nk_angulars_symmetric_bf16_best | nk_bf16_t | nk_f32_t | 24 |
nk_angulars_symmetric_e4m3_best | nk_e4m3_t | nk_f32_t | 22 |
nk_angulars_symmetric_e5m2_best | nk_e5m2_t | nk_f32_t | 23 |
nk_angulars_symmetric_e2m3_best | nk_e2m3_t | nk_f32_t | 24 |
nk_angulars_symmetric_e2m1_best | nk_e2m1x2_t | nk_f32_t | 24 |
nk_angulars_symmetric_nvfp4_best | nk_nvfp4_cref_t | nk_f32_t | 11 |
nk_angulars_symmetric_mxfp4_best | nk_mxfp4_cref_t | nk_f32_t | 11 |
nk_angulars_symmetric_mxfp6e2m3_best | nk_mxfp6e2m3_cref_t | nk_f32_t | 8 |
nk_angulars_symmetric_mxfp6e3m2_best | nk_mxfp6e3m2_cref_t | nk_f32_t | 7 |
nk_angulars_symmetric_mxfp8e4m3_best | nk_mxfp8e4m3_cref_t | nk_f32_t | 11 |
nk_angulars_symmetric_mxfp8e5m2_best | nk_mxfp8e5m2_cref_t | nk_f32_t | 11 |
nk_angulars_symmetric_e3m2_best | nk_e3m2_t | nk_f32_t | 20 |
nk_angulars_symmetric_i8_best | nk_i8_t | nk_f32_t | 23 |
nk_angulars_symmetric_u8_best | nk_u8_t | nk_f32_t | 23 |
nk_angulars_symmetric_i4_best | nk_i4x2_t | nk_f32_t | 16 |
nk_angulars_symmetric_u4_best | nk_u4x2_t | nk_f32_t | 16 |
Backends
nk_angulars_symmetric_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_angulars_symmetric_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_angulars_symmetric_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_angulars_symmetric_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_angulars_symmetric_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_angulars_symmetric_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_angulars_symmetric_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_angulars_symmetric_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_angulars_symmetric_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_angulars_symmetric_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_symmetric_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_symmetric_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_angulars_symmetric_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10
Parameters
vectors | in | Input matrix of row vectors in row-major order. |
|---|---|---|
vector_count | in | Number of vectors (rows) in the input matrix. |
depth | in | Counts dimensions, a multiple of the values per byte. |
stride | in | Row stride in bytes for the input matrix. |
result | out | Output symmetric matrix of vector_count × vector_count. |
result_stride | in | Row stride in bytes for the result matrix. |
rows_begin | in | First row of results to compute (for parallelism). |
rows_end | in | One past the last row of results to compute, clamped to vector_count. |
workspace | in | Scratch memory aligned to nk_default_alignment_k, in device memory ordered by stream on GPU capabilities. May be null with workspace_bytes of 0. |
workspace_bytes | in | At least what the matching workspace-size query reports for the fast path, or 0 for the scratch-free path. Anything in between is nk_unexpected_dimensions_k. |
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_misaligned_kA misalignedworkspace.nk_missing_kernel_kNo capability incapabilitieshas it.
nk_euclideans_packed_*_best× 19Computes batched euclidean distances using a packed second matrix.
nk_status_t nk_euclideans_packed_f32_best(
nk_f32_t const *a,
void const *b_packed,
nk_f64_t *result,
nk_size_t rows,
nk_size_t columns,
nk_size_t depth,
nk_size_t a_stride,
nk_size_t r_stride,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | a | result | Backends |
|---|---|---|---|
nk_euclideans_packed_f32_best | nk_f32_t | nk_f64_t | 12 |
nk_euclideans_packed_f64_best | nk_f64_t | nk_f64_t | 12 |
nk_euclideans_packed_f16_best | nk_f16_t | nk_f32_t | 22 |
nk_euclideans_packed_bf16_best | nk_bf16_t | nk_f32_t | 24 |
nk_euclideans_packed_e4m3_best | nk_e4m3_t | nk_f32_t | 22 |
nk_euclideans_packed_e5m2_best | nk_e5m2_t | nk_f32_t | 23 |
nk_euclideans_packed_e2m3_best | nk_e2m3_t | nk_f32_t | 24 |
nk_euclideans_packed_e2m1_best | nk_e2m1x2_t | nk_f32_t | 24 |
nk_euclideans_packed_nvfp4_best | nk_nvfp4_cref_t | nk_f32_t | 11 |
nk_euclideans_packed_mxfp4_best | nk_mxfp4_cref_t | nk_f32_t | 11 |
nk_euclideans_packed_mxfp6e2m3_best | nk_mxfp6e2m3_cref_t | nk_f32_t | 8 |
nk_euclideans_packed_mxfp6e3m2_best | nk_mxfp6e3m2_cref_t | nk_f32_t | 7 |
nk_euclideans_packed_mxfp8e4m3_best | nk_mxfp8e4m3_cref_t | nk_f32_t | 11 |
nk_euclideans_packed_mxfp8e5m2_best | nk_mxfp8e5m2_cref_t | nk_f32_t | 11 |
nk_euclideans_packed_e3m2_best | nk_e3m2_t | nk_f32_t | 20 |
nk_euclideans_packed_i8_best | nk_i8_t | nk_f32_t | 23 |
nk_euclideans_packed_u8_best | nk_u8_t | nk_f32_t | 23 |
nk_euclideans_packed_i4_best | nk_i4x2_t | nk_f32_t | 16 |
nk_euclideans_packed_u4_best | nk_u4x2_t | nk_f32_t | 16 |
Backends
nk_euclideans_packed_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_euclideans_packed_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_euclideans_packed_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_euclideans_packed_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_euclideans_packed_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_euclideans_packed_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_euclideans_packed_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_euclideans_packed_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_euclideans_packed_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_packed_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_packed_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_packed_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_packed_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10
Parameters
a | in | Input A matrix in row-major order. |
|---|---|---|
b_packed | in | Packed B with norms from nk_dots_pack_*_best on the same capabilities. |
result | out | Output matrix (rows x columns) of euclidean distances. |
rows | in | Number of rows in A. |
columns | in | Number of columns in B (packed). |
depth | in | Shared inner dimension in dimensions, a multiple of the values per byte. |
a_stride | in | Row stride in bytes for A. |
r_stride | in | Row stride in bytes for the result matrix. |
workspace | in | Scratch memory aligned to nk_default_alignment_k, in device memory ordered by stream on GPU capabilities. May be null with workspace_bytes of 0. |
workspace_bytes | in | At least what the matching workspace-size query reports for the fast path, or 0 for the scratch-free path. Anything in between is nk_unexpected_dimensions_k. |
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_misaligned_kA misalignedworkspace.nk_missing_kernel_kNo capability incapabilitieshas it.
nk_euclideans_symmetric_*_best× 19Computes symmetric euclidean distance matrix (Gram-style) for a set of vectors.
nk_status_t nk_euclideans_symmetric_f32_best(
nk_f32_t const *vectors,
nk_size_t vector_count,
nk_size_t depth,
nk_size_t stride,
nk_f64_t *result,
nk_size_t result_stride,
nk_size_t rows_begin,
nk_size_t rows_end,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | vectors | result | Backends |
|---|---|---|---|
nk_euclideans_symmetric_f32_best | nk_f32_t | nk_f64_t | 12 |
nk_euclideans_symmetric_f64_best | nk_f64_t | nk_f64_t | 12 |
nk_euclideans_symmetric_f16_best | nk_f16_t | nk_f32_t | 22 |
nk_euclideans_symmetric_bf16_best | nk_bf16_t | nk_f32_t | 24 |
nk_euclideans_symmetric_e4m3_best | nk_e4m3_t | nk_f32_t | 22 |
nk_euclideans_symmetric_e5m2_best | nk_e5m2_t | nk_f32_t | 23 |
nk_euclideans_symmetric_e2m3_best | nk_e2m3_t | nk_f32_t | 24 |
nk_euclideans_symmetric_e2m1_best | nk_e2m1x2_t | nk_f32_t | 24 |
nk_euclideans_symmetric_nvfp4_best | nk_nvfp4_cref_t | nk_f32_t | 11 |
nk_euclideans_symmetric_mxfp4_best | nk_mxfp4_cref_t | nk_f32_t | 11 |
nk_euclideans_symmetric_mxfp6e2m3_best | nk_mxfp6e2m3_cref_t | nk_f32_t | 8 |
nk_euclideans_symmetric_mxfp6e3m2_best | nk_mxfp6e3m2_cref_t | nk_f32_t | 7 |
nk_euclideans_symmetric_mxfp8e4m3_best | nk_mxfp8e4m3_cref_t | nk_f32_t | 11 |
nk_euclideans_symmetric_mxfp8e5m2_best | nk_mxfp8e5m2_cref_t | nk_f32_t | 11 |
nk_euclideans_symmetric_e3m2_best | nk_e3m2_t | nk_f32_t | 20 |
nk_euclideans_symmetric_i8_best | nk_i8_t | nk_f32_t | 23 |
nk_euclideans_symmetric_u8_best | nk_u8_t | nk_f32_t | 23 |
nk_euclideans_symmetric_i4_best | nk_i4x2_t | nk_f32_t | 16 |
nk_euclideans_symmetric_u4_best | nk_u4x2_t | nk_f32_t | 16 |
Backends
nk_euclideans_symmetric_f32serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_euclideans_symmetric_f64serialsmef64haswellskylakev128relaxedneonrvvpowervsxloongsonasxcudaampererocmnk_euclideans_symmetric_f16serialgraniteamxsmehaswellskylakev128relaxedneonneonfhmrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_bf16serialgenoasapphireamxsmehaswellskylakev128relaxedv128neonneonbfdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_e4m3serialgenoadiamondsapphireamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_e5m2serialgenoadiamondsapphireamxgraniteamxsmehaswellskylakev128relaxedneonfhmneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_e2m3serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_e2m1serialsapphireamxsmehaswellskylakeicelakealdersierrav128relaxedneonsdotneonfp8rvvcudaamperehopperblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_nvfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_euclideans_symmetric_mxfp4serialsapphireamxsmeskylakeneonneonsdotcudablackwellmetalapple9apple10nk_euclideans_symmetric_mxfp6e2m3serialsmeskylakeneonneonsdotmetalapple9apple10nk_euclideans_symmetric_mxfp6e3m2serialsmeskylakeneonmetalapple9apple10nk_euclideans_symmetric_mxfp8e4m3serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_euclideans_symmetric_mxfp8e5m2serialsapphireamxsmeskylakeneonneonfhmcudablackwellmetalapple9apple10nk_euclideans_symmetric_e3m2serialsapphireamxsmehaswellskylakev128relaxedneonsdotneonfp8rvvcudaampereblackwellblackwellrtxrocmcdna3cdna4cdna5metalapple9apple10nk_euclideans_symmetric_i8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_symmetric_u8serialsapphireamxsmehaswellicelakealdersierrav128relaxedv128neonsdotrvvpowervsxloongsonasxcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_symmetric_i4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10nk_euclideans_symmetric_u4serialsmehaswellicelakev128relaxedneonsdotcudaamperehopperblackwellrocmcdna3cdna4cdna5metalapple10
Parameters
vectors | in | Input matrix of row vectors in row-major order. |
|---|---|---|
vector_count | in | Number of vectors (rows) in the input matrix. |
depth | in | Counts dimensions, a multiple of the values per byte. |
stride | in | Row stride in bytes for the input matrix. |
result | out | Output symmetric matrix of vector_count × vector_count. |
result_stride | in | Row stride in bytes for the result matrix. |
rows_begin | in | First row of results to compute (for parallelism). |
rows_end | in | One past the last row of results to compute, clamped to vector_count. |
workspace | in | Scratch memory aligned to nk_default_alignment_k, in device memory ordered by stream on GPU capabilities. May be null with workspace_bytes of 0. |
workspace_bytes | in | At least what the matching workspace-size query reports for the fast path, or 0 for the scratch-free path. Anything in between is nk_unexpected_dimensions_k. |
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_misaligned_kA misalignedworkspace.nk_missing_kernel_kNo capability incapabilitieshas it.
nk_attention_pack_size_*_best× 10Returns bytes enough to pack a ragged batch of segments however its tokens split.
nk_status_t nk_attention_pack_size_bf16_best(
nk_size_t key_value_head_count,
nk_size_t depth,
nk_size_t token_count,
nk_size_t segment_count,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_attention_pack_size_bf16_best | 22 |
nk_attention_pack_size_f16_best | 11 |
nk_attention_pack_size_e4m3_best | 23 |
nk_attention_pack_size_i8_best | 20 |
nk_attention_pack_size_nvfp4_best | 7 |
nk_attention_pack_size_mxfp4_best | 7 |
nk_attention_pack_size_mxfp6e2m3_best | 7 |
nk_attention_pack_size_mxfp6e3m2_best | 7 |
nk_attention_pack_size_mxfp8e4m3_best | 7 |
nk_attention_pack_size_mxfp8e5m2_best | 7 |
Backends
nk_attention_pack_size_bf16serialhaswellskylakegenoasapphireamxsmeneonneonbfdotrvvv128relaxedcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple9apple10nk_attention_pack_size_f16serialsmeneonneonfhmcudaampereblackwellblackwellultrametalapple9apple10nk_attention_pack_size_e4m3serialhaswellskylakegenoasapphireamxdiamondamxsmeneonneonfhmrvvv128relaxedcudaamperehopperblackwellblackwellrtxblackwellultrarocmcdna4cdna5metalapple9apple10nk_attention_pack_size_i8serialhaswellicelakesapphireamxsmeneonneonsdotrvvv128relaxedcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple10nk_attention_pack_size_nvfp4serialsmeneonneonfhmmetalapple9apple10nk_attention_pack_size_mxfp4serialsmeneonneonbfdotmetalapple9apple10nk_attention_pack_size_mxfp6e2m3serialsmeneonneonbfdotmetalapple9apple10nk_attention_pack_size_mxfp6e3m2serialsmeneonneonbfdotmetalapple9apple10nk_attention_pack_size_mxfp8e4m3serialsmeneonneonbfdotmetalapple9apple10nk_attention_pack_size_mxfp8e5m2serialsmeneonneonbfdotmetalapple9apple10
Parameters
key_value_head_count | in | Number of K/V heads, a nonzero divisor of the query head count; attention must take it and depth as packed, and only debug builds assert these rules. |
|---|---|---|
depth | in | Head dimension; any value ≥ 1. |
token_count | in | Tokens across all segments, the sum of the pack's key counts. |
segment_count | in | Number of segments packed together. |
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_kNo capability incapabilitieshas it.
nk_attention_packed_shape_*_best× 10Reads a packed KV cache's shape from its header.
nk_status_t nk_attention_packed_shape_bf16_best(
void const *key_value_packed,
nk_size_t *key_value_head_count,
nk_size_t *depth,
nk_size_t *segments,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | Backends |
|---|---|
nk_attention_packed_shape_bf16_best | 22 |
nk_attention_packed_shape_f16_best | 11 |
nk_attention_packed_shape_e4m3_best | 23 |
nk_attention_packed_shape_i8_best | 20 |
nk_attention_packed_shape_nvfp4_best | 7 |
nk_attention_packed_shape_mxfp4_best | 7 |
nk_attention_packed_shape_mxfp6e2m3_best | 7 |
nk_attention_packed_shape_mxfp6e3m2_best | 7 |
nk_attention_packed_shape_mxfp8e4m3_best | 7 |
nk_attention_packed_shape_mxfp8e5m2_best | 7 |
Backends
nk_attention_packed_shape_bf16serialhaswellskylakegenoasapphireamxsmeneonneonbfdotrvvv128relaxedcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple9apple10nk_attention_packed_shape_f16serialsmeneonneonfhmcudaampereblackwellblackwellultrametalapple9apple10nk_attention_packed_shape_e4m3serialhaswellskylakegenoasapphireamxdiamondamxsmeneonneonfhmrvvv128relaxedcudaamperehopperblackwellblackwellrtxblackwellultrarocmcdna4cdna5metalapple9apple10nk_attention_packed_shape_i8serialhaswellicelakesapphireamxsmeneonneonsdotrvvv128relaxedcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple10nk_attention_packed_shape_nvfp4serialsmeneonneonfhmmetalapple9apple10nk_attention_packed_shape_mxfp4serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_shape_mxfp6e2m3serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_shape_mxfp6e3m2serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_shape_mxfp8e4m3serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_shape_mxfp8e5m2serialsmeneonneonbfdotmetalapple9apple10
Parameters
key_value_packed | in | A buffer produced by the matching nk_attention_pack_bf16_best. |
|---|---|---|
key_value_head_count | out | Receives the K/V head count. |
depth | out | Receives the head dimension. |
segments | out | Receives the segment count. |
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_kNo capability incapabilitieshas it.nk_device_code_mismatch_kA GPU stream capturing a graph, which reading the header back would void.
nk_attention_packed_segmentsLocates a pack's key slot boundaries and key counts.
nk_status_t nk_attention_packed_segments(
void const *key_value_packed,
nk_size_t segment_count,
nk_u32_t const **key_offsets,
nk_u32_t const **key_lengths);Parameters
key_value_packed | in | A buffer produced by the matching nk_attention_pack_bf16_best. |
|---|---|---|
segment_count | in | The pack's segment count, from nk_attention_packed_shape_bf16_best. |
key_offsets | out | Points at the pack's segment_count + 1 slot boundaries. |
key_lengths | out | Points at the pack's segment_count key counts. |
Notes
Pure address arithmetic that never dereferences the pack, so it works on device packs too.
nk_attention_pack_*_best× 4Packs a ragged batch of K and V segments into a backend-opaque layout.
nk_status_t nk_attention_pack_bf16_best(
nk_size_t key_value_head_count,
nk_size_t depth,
nk_u32_t const *key_offsets,
nk_u32_t const *key_lengths,
nk_size_t segment_count,
nk_bf16_t const *keys,
nk_size_t key_stride,
nk_bf16_t const *values,
nk_size_t value_stride,
void *key_value_packed,
nk_size_t tasks_begin,
nk_size_t tasks_end,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | keys, values | Backends |
|---|---|---|
nk_attention_pack_bf16_best | nk_bf16_t | 22 |
nk_attention_pack_f16_best | nk_f16_t | 11 |
nk_attention_pack_e4m3_best | nk_e4m3_t | 23 |
nk_attention_pack_i8_best | nk_i8_t | 20 |
Backends
nk_attention_pack_bf16serialhaswellskylakegenoasapphireamxsmeneonneonbfdotrvvv128relaxedcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple9apple10nk_attention_pack_f16serialsmeneonneonfhmcudaampereblackwellblackwellultrametalapple9apple10nk_attention_pack_e4m3serialhaswellskylakegenoasapphireamxdiamondamxsmeneonneonfhmrvvv128relaxedcudaamperehopperblackwellblackwellrtxblackwellultrarocmcdna4cdna5metalapple9apple10nk_attention_pack_i8serialhaswellicelakesapphireamxsmeneonneonsdotrvvv128relaxedcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple10
GPU packs read the arrays on the device, so only their counts are checked. Windows let callers pack in parallel and in any order: tasks write disjoint ranges, the window starting at task 0 also writes the header and directory, and no window reads them.
Parameters
key_value_head_count | in | K/V heads, at most the query head count. |
|---|---|---|
depth | in | Head dimension. |
key_offsets | in | The segment_count + 1 slot boundaries: segment s keeps its keys in rows key_offsets[s] onward, and may leave spare capacity before key_offsets[s + 1]. 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: the difference of consecutive key_offsets. |
segment_count | in | Number of segments in the batch. |
keys | in | Token-major matrix, one row of key_value_head_count × depth contiguous elements per token, with strided rows. Packing fuses no transposition, so depth-major K is transposed in a separate pass first. |
key_stride | in | Row (token) stride of keys in bytes. |
values | in | Token-major matrix shaped like keys. Depth-major V is likewise transposed in a separate pass first. |
value_stride | in | Row (token) stride of values in bytes. |
key_value_packed | out | 64-byte-aligned buffer of nk_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. |
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_unexpected_dimensions_kA count overUINT32_MAXor slots that decrease or overflow.nk_missing_kernel_kNo capability incapabilitieshas it.
nk_attention_pack_*_best× 6Packs block-scaled K and V as nk_attention_pack_bf16_best does, from references to their codes, block scales and NVFP4 tensor scale; depth is a whole number of blocks, and the strides count bytes of codes.
nk_status_t nk_attention_pack_nvfp4_best(
nk_size_t key_value_head_count,
nk_size_t depth,
nk_u32_t const *key_offsets,
nk_u32_t const *key_lengths,
nk_size_t segment_count,
nk_nvfp4_cref_t const *keys,
nk_size_t key_stride,
nk_nvfp4_cref_t const *values,
nk_size_t value_stride,
void *key_value_packed,
nk_size_t tasks_begin,
nk_size_t tasks_end,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | keys, values | Backends |
|---|---|---|
nk_attention_pack_nvfp4_best | nk_nvfp4_cref_t | 7 |
nk_attention_pack_mxfp4_best | nk_mxfp4_cref_t | 7 |
nk_attention_pack_mxfp6e2m3_best | nk_mxfp6e2m3_cref_t | 7 |
nk_attention_pack_mxfp6e3m2_best | nk_mxfp6e3m2_cref_t | 7 |
nk_attention_pack_mxfp8e4m3_best | nk_mxfp8e4m3_cref_t | 7 |
nk_attention_pack_mxfp8e5m2_best | nk_mxfp8e5m2_cref_t | 7 |
Backends
nk_attention_pack_nvfp4serialsmeneonneonfhmmetalapple9apple10nk_attention_pack_mxfp4serialsmeneonneonbfdotmetalapple9apple10nk_attention_pack_mxfp6e2m3serialsmeneonneonbfdotmetalapple9apple10nk_attention_pack_mxfp6e3m2serialsmeneonneonbfdotmetalapple9apple10nk_attention_pack_mxfp8e4m3serialsmeneonneonbfdotmetalapple9apple10nk_attention_pack_mxfp8e5m2serialsmeneonneonbfdotmetalapple9apple10
nk_attention_packed_*_best× 10Ragged scaled-dot-product attention under a band of visible keys.
nk_status_t nk_attention_packed_bf16_best(
nk_size_t head_count,
nk_size_t key_value_head_count,
nk_size_t depth,
nk_u32_t const *query_offsets,
nk_size_t query_token_count,
nk_f32_t scale,
nk_size_t keys_before,
nk_size_t keys_after,
nk_bf16_t const *queries,
nk_size_t query_stride,
void const *key_value_packed,
nk_f32_t *output,
nk_size_t output_stride,
nk_f32_t *log_sum_exp,
nk_size_t tasks_begin,
nk_size_t tasks_end,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | queries | Backends |
|---|---|---|
nk_attention_packed_bf16_best | nk_bf16_t | 22 |
nk_attention_packed_f16_best | nk_f16_t | 11 |
nk_attention_packed_e4m3_best | nk_e4m3_t | 23 |
nk_attention_packed_i8_best | nk_i8_t | 20 |
nk_attention_packed_nvfp4_best | nk_nvfp4_cref_t | 7 |
nk_attention_packed_mxfp4_best | nk_mxfp4_cref_t | 7 |
nk_attention_packed_mxfp6e2m3_best | nk_mxfp6e2m3_cref_t | 7 |
nk_attention_packed_mxfp6e3m2_best | nk_mxfp6e3m2_cref_t | 7 |
nk_attention_packed_mxfp8e4m3_best | nk_mxfp8e4m3_cref_t | 7 |
nk_attention_packed_mxfp8e5m2_best | nk_mxfp8e5m2_cref_t | 7 |
Backends
nk_attention_packed_bf16serialhaswellskylakegenoasapphireamxsmeneonneonbfdotrvvv128relaxedcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple9apple10nk_attention_packed_f16serialsmeneonneonfhmcudaampereblackwellblackwellultrametalapple9apple10nk_attention_packed_e4m3serialhaswellskylakegenoasapphireamxdiamondamxsmeneonneonfhmrvvv128relaxedcudaamperehopperblackwellblackwellrtxblackwellultrarocmcdna4cdna5metalapple9apple10nk_attention_packed_i8serialhaswellicelakesapphireamxsmeneonneonsdotrvvv128relaxedcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple10nk_attention_packed_nvfp4serialsmeneonneonfhmmetalapple9apple10nk_attention_packed_mxfp4serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_mxfp6e2m3serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_mxfp6e3m2serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_mxfp8e4m3serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_mxfp8e5m2serialsmeneonneonbfdotmetalapple9apple10
Computes O[s] = softmax(Q[s] × K[s]ᵀ × scale) × V[s] for every segment s over the keys each query row sees, as Causal, Sliding-Window and Bidirectional Masks defines. Covers self-attention, where query_offsets equal the pack-time key_offsets, cross-attention, and pooling, where query_offsets = {0, 1, 2, …} holds one query per segment, plus GQA/MQA via key_value_head_count < head_count. Pack key_value_packed with the same capabilities: CPU kernels return nk_pack_mismatch_k for buffers another capability packed, while GPU kernels trust it, as reading its device header would wait on the stream.
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: NUMKONG_SIZE_MAX for causal and bidirectional attention, w − 1 for a sliding window of w keys. |
keys_after | in | Keys visible after each query's position: 0 for causal attention, NUMKONG_SIZE_MAX for bidirectional. |
queries | in | Token-major matrix, one row of head_count × depth elements per query token, with query_stride bytes between rows. |
query_stride | in | Row (token) stride of queries in bytes. |
key_value_packed | in | Buffer produced by the matching nk_attention_pack_* backend. |
output | out | Token-major F32 matrix, one row of head_count × depth elements per query token, with output_stride bytes between rows. |
output_stride | in | Row (token) stride of output in bytes. |
log_sum_exp | out | Optional natural log of Σ exp(score multiplier × q · k) over each row's keys, one per query token and head, [query_tokens,heads] with no padding, −∞ for a row without keys: what the backward pass recomputes the weights from. Null skips it. |
tasks_begin | in | First task of a window over the query tokens × heads grid, task t × head_count + h being head h of query token t, as in log_sum_exp. |
tasks_end | in | End of that half-open window, clipped to the grid's live rows from query_offsets[0] × head_count. A window partitions the rows of output and log_sum_exp the call writes, per query token × head, so callers parallelize freely; a window may split a segment or a head group. |
workspace | in | Scratch memory on nk_default_alignment_k, on the device of stream, used in its order, so the caller frees it only once the stream passes the call. |
workspace_bytes | in | At least what nk_attention_packed_workspace_size_bf16_best reports for the fast path, or 0 for the path without scratch. |
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_misaligned_kA misalignedworkspace.nk_unexpected_dimensions_kworkspace_bytesbetween 0 and the reported size.nk_missing_kernel_kNo capability incapabilitieshas it.
nk_attention_packed_workspace_size_*_best× 10Returns the workspace bytes nk_attention_packed_bf16_best takes for its fast path on these shapes, by host arithmetic alone.
nk_status_t nk_attention_packed_workspace_size_bf16_best(
nk_size_t head_count,
nk_size_t key_value_head_count,
nk_size_t depth,
nk_size_t query_token_count,
nk_size_t keys_before,
nk_size_t keys_after,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_attention_packed_workspace_size_bf16_best | 22 |
nk_attention_packed_workspace_size_f16_best | 11 |
nk_attention_packed_workspace_size_e4m3_best | 23 |
nk_attention_packed_workspace_size_i8_best | 20 |
nk_attention_packed_workspace_size_nvfp4_best | 7 |
nk_attention_packed_workspace_size_mxfp4_best | 7 |
nk_attention_packed_workspace_size_mxfp6e2m3_best | 7 |
nk_attention_packed_workspace_size_mxfp6e3m2_best | 7 |
nk_attention_packed_workspace_size_mxfp8e4m3_best | 7 |
nk_attention_packed_workspace_size_mxfp8e5m2_best | 7 |
Backends
nk_attention_packed_workspace_size_bf16serialhaswellskylakegenoasapphireamxsmeneonneonbfdotrvvv128relaxedcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple9apple10nk_attention_packed_workspace_size_f16serialsmeneonneonfhmcudaampereblackwellblackwellultrametalapple9apple10nk_attention_packed_workspace_size_e4m3serialhaswellskylakegenoasapphireamxdiamondamxsmeneonneonfhmrvvv128relaxedcudaamperehopperblackwellblackwellrtxblackwellultrarocmcdna4cdna5metalapple9apple10nk_attention_packed_workspace_size_i8serialhaswellicelakesapphireamxsmeneonneonsdotrvvv128relaxedcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple10nk_attention_packed_workspace_size_nvfp4serialsmeneonneonfhmmetalapple9apple10nk_attention_packed_workspace_size_mxfp4serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_workspace_size_mxfp6e2m3serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_workspace_size_mxfp6e3m2serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_workspace_size_mxfp8e4m3serialsmeneonneonbfdotmetalapple9apple10nk_attention_packed_workspace_size_mxfp8e5m2serialsmeneonneonbfdotmetalapple9apple10
Only the Metal capabilities need one, for decoding at most four query rows against long keys, so every other capability reports 0. The bytes hold whenever the shapes allow the fast path, which may still decline it for the call's pointers or key count.
Every other parameter follows nk_attention_packed_bf16_best.
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. |
keys_after | in | Keys visible after each query's position. |
capabilities | in | One device's capabilities, like nk_cpu_capabilities_enabled reports. |
bytes | out | The workspace size, 0 when the capability needs none. |
Statuses
nk_missing_kernel_kNo capability incapabilitieshas it.
nk_attention_packed_gradients_*_best× 7Gradients of ragged attention with respect to its queries, keys and values.
nk_status_t nk_attention_packed_gradients_bf16_best(
nk_size_t head_count,
nk_size_t key_value_head_count,
nk_size_t depth,
nk_u32_t const *query_offsets,
nk_size_t query_token_count,
nk_f32_t scale,
nk_size_t keys_before,
nk_size_t keys_after,
nk_bf16_t const *queries,
nk_size_t query_stride,
void const *key_value_packed,
nk_f32_t const *output,
nk_f32_t const *output_gradient,
nk_size_t output_stride,
nk_f32_t const *log_sum_exp,
nk_f32_t *query_gradient,
nk_size_t query_gradient_stride,
nk_f32_t *key_gradient,
nk_f32_t *value_gradient,
nk_size_t key_value_gradient_stride,
nk_size_t tasks_begin,
nk_size_t tasks_end,
void *workspace,
nk_size_t workspace_bytes,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | queries | Backends |
|---|---|---|
nk_attention_packed_gradients_bf16_best | nk_bf16_t | 13 |
nk_attention_packed_gradients_nvfp4_best | nk_nvfp4_cref_t | 1 |
nk_attention_packed_gradients_mxfp4_best | nk_mxfp4_cref_t | 1 |
nk_attention_packed_gradients_mxfp6e2m3_best | nk_mxfp6e2m3_cref_t | 1 |
nk_attention_packed_gradients_mxfp6e3m2_best | nk_mxfp6e3m2_cref_t | 1 |
nk_attention_packed_gradients_mxfp8e4m3_best | nk_mxfp8e4m3_cref_t | 1 |
nk_attention_packed_gradients_mxfp8e5m2_best | nk_mxfp8e5m2_cref_t | 1 |
Backends
nk_attention_packed_gradients_bf16serialcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple9apple10nk_attention_packed_gradients_nvfp4serialnk_attention_packed_gradients_mxfp4serialnk_attention_packed_gradients_mxfp6e2m3serialnk_attention_packed_gradients_mxfp6e3m2serialnk_attention_packed_gradients_mxfp8e4m3serialnk_attention_packed_gradients_mxfp8e5m2serial
Recomputes every attention weight from log_sum_exp, so nothing of the forward but its output and that per-row statistic is kept. The gradients are F32 and overwritten: query_gradient holds one row of head_count × depth elements per query token, while key_gradient and value_gradient hold one row of key_value_head_count × depth elements per key token, token t of segment s at row key_offsets[s] + t, with the offsets the pack stored. Rows past a segment's key count, the gaps and tails before the next segment's first key, are never written, so callers zero them when they matter. Tasks cover the segments × key-value heads grid, each owning one KV head of one segment together with the query heads sharing it, so tasks write disjoint gradients and each sums in a fixed order. The Blackwell kernels run a single pass for BF16 heads of 64 or 128 dimensions with 16-byte aligned rows, given a workspace of query_token_count × head_count × (2 × depth + 4) bytes, as nk_attention_packed_gradients_workspace_size_bf16_best reports, and add the query_gradient partials of different key blocks in F32 in any order, so their low bits may vary run to run. Other calls, and calls without a workspace, take the two-pass path.
Every other parameter follows nk_attention_packed_bf16_best.
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, with query_stride bytes between token rows. |
query_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 natural log of Σ exp(scale × score) over each row's keys, one per query token and head, [query_tokens,heads] with no padding. |
query_gradient | out | Gradient with respect to queries. |
query_gradient_stride | in | Bytes between the query_gradient rows of consecutive tokens. |
key_gradient | out | Gradient with respect to the packed keys. |
value_gradient | out | Gradient with respect to the packed values. |
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. A 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 on nk_default_alignment_k, on the device of stream. |
workspace_bytes | in | At least what nk_attention_packed_gradients_workspace_size_bf16_best reports for the fast path, or 0 for the path without scratch. |
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. |
nk_attention_packed_gradients_workspace_size_*_best× 7Returns the workspace bytes nk_attention_packed_gradients_bf16_best takes for its fast path on these shapes, as nk_attention_packed_workspace_size_bf16_best does for attention: the Blackwell single pass for BF16 is the only one that needs any.
nk_status_t nk_attention_packed_gradients_workspace_size_bf16_best(
nk_size_t head_count,
nk_size_t key_value_head_count,
nk_size_t depth,
nk_size_t query_token_count,
nk_size_t keys_before,
nk_size_t keys_after,
nk_capability_t capabilities,
nk_size_t *bytes);| Symbol | Backends |
|---|---|
nk_attention_packed_gradients_workspace_size_bf16_best | 13 |
nk_attention_packed_gradients_workspace_size_nvfp4_best | 1 |
nk_attention_packed_gradients_workspace_size_mxfp4_best | 1 |
nk_attention_packed_gradients_workspace_size_mxfp6e2m3_best | 1 |
nk_attention_packed_gradients_workspace_size_mxfp6e3m2_best | 1 |
nk_attention_packed_gradients_workspace_size_mxfp8e4m3_best | 1 |
nk_attention_packed_gradients_workspace_size_mxfp8e5m2_best | 1 |
Backends
nk_attention_packed_gradients_workspace_size_bf16serialcudaamperehopperblackwellblackwellultrarocmcdna3cdna4cdna5metalapple9apple10nk_attention_packed_gradients_workspace_size_nvfp4serialnk_attention_packed_gradients_workspace_size_mxfp4serialnk_attention_packed_gradients_workspace_size_mxfp6e2m3serialnk_attention_packed_gradients_workspace_size_mxfp6e3m2serialnk_attention_packed_gradients_workspace_size_mxfp8e4m3serialnk_attention_packed_gradients_workspace_size_mxfp8e5m2serial
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 | One device's capabilities, like nk_cpu_capabilities_enabled reports. |
bytes | out | The workspace size, 0 when the capability needs none. |
nk_attention_rope_*_best× 3NeoX split-half rotary position embedding (RoPE): rotates channel pairs by per-token angles.
nk_status_t nk_attention_rope_f32_best(
nk_f32_t const *x,
nk_f32_t const *cos,
nk_f32_t const *sin,
nk_f32_t *y,
nk_size_t rows,
nk_size_t head_count,
nk_size_t depth,
nk_size_t x_stride,
nk_size_t y_stride,
nk_capability_t capabilities,
nk_stream_t stream);| Symbol | x, y | Backends |
|---|---|---|
nk_attention_rope_f32_best | nk_f32_t | 6 |
nk_attention_rope_bf16_best | nk_bf16_t | 6 |
nk_attention_rope_e4m3_best | nk_e4m3_t | 6 |
Backends
nk_attention_rope_f32serialhaswellskylakeneoncudametalnk_attention_rope_bf16serialhaswellskylakeneoncudametalnk_attention_rope_e4m3serialhaswellskylakeneoncudametal
Rotates the channel pair i and i + h of every head in every row, h being half the depth:
* y[i] = x[i] · cos - x[i + h] · sin
* y[i + h] = x[i] · sin + x[i + h] · cos
* The rotation is linear, so FP8 codes rotate at their own scale and need none passed.
Parameters
x | in | Input token matrix of shape rows by head_count × depth. |
|---|---|---|
cos | in | Per-token cosine angle grid of shape rows by depth / 2, shared by heads. |
sin | in | Per-token sine angle grid of shape rows by depth / 2, shared across heads. |
y | out | Output matrix, same shape and dtype as x; may alias x for in-place rotation. |
rows | in | The number of token rows. |
head_count | in | The number of heads per token. |
depth | in | The even number of channels per head. |
x_stride | in | Row stride of x in bytes. |
y_stride | in | Row stride of y in bytes. |
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_kNo capability incapabilitieshas it.
nk_attention_output_dtypeReturns the output dtype for attention: accumulator-precision F32 for all inputs.
nk_dtype_t nk_attention_output_dtype(nk_dtype_t dtype);nk_attention_rope_error_boundReturns the error bound of RoPE outputs before they round into their type, relative to the two products each of them sums: 3 roundings in F32, which every input type rotates in, through the descale, the product and the sum.
nk_f64_t nk_attention_rope_error_bound(nk_dtype_t dtype);