Matrix Kernels

Packed matrix kernels are the right tool when the right-hand side is reused across many query batches. This is the GEMM-like story.

import numpy as np
import numkong as nk

left = np.random.randn(128, 768).astype(np.float32)
right = np.random.randn(10_000, 768).astype(np.float32)

right_packed = nk.dots_pack(right, dtype="float32")  # pack once, reuse many times
scores = nk.dots_packed(left, right_packed)          # equivalent to left @ right.T

assert scores.shape == (128, 10_000)
assert right_packed.nbytes == nk.PackedMatrix.pack_size(10_000, 768, dtype="float32")

Important runtime rules from the current implementation:

• a must be rank-2
• a must have contiguous rows
• negative strides are rejected for these matrix kernels
• out, when provided, must be C-contiguous with the expected dtype
• rows_begin and rows_end split the left operand rows

The arithmetic advantages are:

• one-time packing of B
• one-time internal layout conversion and depth padding
• norm reuse for angulars_packed and euclideans_packed
• no repeated scan of the original right-hand-side layout

Packing itself does not require aligned caller buffers. The packed object owns its internal payload and handles the layout under the hood.

Tensor @ PackedMatrix is also supported and maps to the same packed dot-product path.

§ All-Pairs APIs and cdist

cdist is the NumPy/SciPy-shaped all-pairs entrypoint. It handles rectangular matrix pairs and symmetric self-distance cases.

import numpy as np
import numkong as nk

queries = np.random.randn(100, 768).astype(np.float32)
database = np.random.randn(10_000, 768).astype(np.float32)

pairwise = nk.angular(queries, database[:100])             # rectangular broadcasted pairwise call
all_pairs = nk.cdist(queries, database, metric="angular")  # scipy.spatial.distance.cdist analogue

assert np.asarray(pairwise).shape == (100, 100)
assert np.asarray(all_pairs).shape == (100, 10_000)

cdist and the packed and symmetric kernels take a threads= argument, where 0 means every logical processor, and row ranges are still available for partitioning the work yourself. dots_pack and hammings_pack take it too, splitting the packed matrix by columns.

§ Symmetric Kernels for SYRK-Like Workloads

Symmetric kernels solve a different problem from packed cross-matrix kernels. They compute self-similarity or self-distance matrices. This is the SYRK-like story.

import numpy as np
import numkong as nk

vectors = np.random.randn(1024, 768).astype(np.float32)
out = nk.zeros((1024, 1024), dtype="float64")

nk.dots_symmetric(vectors, out=out, rows_begin=0, rows_end=256)
nk.dots_symmetric(vectors, out=out, rows_begin=256, rows_end=512)

assert out.shape == (1024, 1024)

This family has different economics from packed GEMM-like work. It avoids duplicate (i, j) and (j, i) evaluations. It is naturally partitioned by row windows of one square output.

angulars_symmetric and euclideans_symmetric also benefit from reuse of dot-product-derived work inside the symmetric sweep. That is why these APIs are faster than a nested Python loop over angular(a[i], a[j]).

§ Attention with a Pre-Packed KV-Cache

Ragged scaled-dot-product attention packs the key/value tokens once, then scores many queries against them. Segments of different lengths share one packed cache, addressed by cumulative key_offsets. A segment may hold fewer keys than its slot, as a decode cache with spare room does: pass key_lengths too, and the keys past each length are left out of the pack.

import numpy as np
import numkong as nk

lengths = [6, 10]  # two ragged segments
offsets = np.array([0, *np.cumsum(lengths)], dtype=np.uint32)
tokens, heads, depth = int(offsets[-1]), 2, 64

k = nk.Tensor(np.random.randn(tokens, heads * depth).astype(np.float32)).astype("bfloat16")
v = nk.Tensor(np.random.randn(tokens, heads * depth).astype(np.float32)).astype("bfloat16")
kv = nk.attention_pack(k, v, key_offsets=offsets, depth=depth)

queries = nk.Tensor(np.random.randn(tokens, heads * depth).astype(np.float32)).astype("bfloat16")
out = nk.attention_packed(queries, kv, query_offsets=offsets)
windowed = nk.attention_packed(queries, kv, query_offsets=offsets, keys_before=3, keys_after=0)

assert kv.shape == (kv.key_value_head_count, kv.depth, kv.segments)

Queries align to the end of each segment’s keys: row r of q queries over k keys sits at position r + k - q and sees keys_before keys before it and keys_after after, None being unbounded. So keys_after=0 is causal for prefill and decode alike, and adding keys_before=w - 1 slides a window of w keys. Rows that see no key come back as zeros.

For training, the forward also fills a (tokens, heads) log-sum-exp, and attention_packed_gradients turns it, the output and its gradient into BF16 query, key and value gradients under the same band. CPUs compute them serially, so the cache is packed with nk.Capability.SERIAL:

kv = nk.attention_pack(k, v, key_offsets=offsets, depth=depth, capabilities=nk.Capability.SERIAL)
log_sum_exp = np.empty((tokens, heads), dtype=np.float32)
out = nk.attention_packed(queries, kv, query_offsets=offsets, keys_after=0, log_sum_exp=log_sum_exp)
output_gradient = np.random.randn(tokens, heads * depth).astype(np.float32)
query_gradient, key_gradient, value_gradient = nk.attention_packed_gradients(
    queries, kv, query_offsets=offsets, output=out, output_gradient=output_gradient,
    log_sum_exp=log_sum_exp, keys_after=0)

query_gradient= takes a pre-allocated (tokens, heads * depth) f32 tensor, whose rows may sit in a wider buffer. With threads=, attention cuts its grid of query tokens × heads into a few windows of equal cost per thread, so a single decode token still spreads over its heads.