# Matrix Kernels

> NumKong for Python

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

```python
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.

```python
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.

```python
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.

```python
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`:

```python
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.

`import numkong`

`from numkong import PackedMatrix, AttentionPackedMatrix`

## PackedMatrix

Opaque pre-packed matrix for repeated cross operations.

```python
class PackedMatrix
```

Created by dots_pack() or hammings_pack() and used with
dots_packed(), hammings_packed(), jaccards_packed(), angulars_packed(),
euclideans_packed(), or the @ operator.

## pack_*_size

Return packed buffer size in bytes for given dimensions, dtype and capabilities, GPU ones included.

```python
@classmethod
def pack__size(
    columns: int,
    depth: int,
    /,
    dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType,
    capabilities: Capability | int | None = None,
) -> int
```

## __repr__*__repr__

Return a string representation.

```python
def __repr____repr__() -> str
```

## columns

Number of packed vectors, the columns of the result.

```python
columns: int
```

## depth

Number of columns in the original matrix.

```python
depth: int
```

## dtype

Data type of the packed matrix (like 'bf16' or 'i8').

```python
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName
```

## nbytes

Size of the packed buffer in bytes.

```python
nbytes: int
```

## shape

Dimensions (columns, depth) read from the packed buffer header.

```python
shape: tuple[int, int]
```

## AttentionPackedMatrix

Opaque pre-packed ragged KV-cache for scaled-dot-product attention.

```python
class AttentionPackedMatrix
```

Created by attention_pack() and used with attention_packed().

## segments

Number of ragged segments.

```python
segments: int
```

## key_value_head_count

Number of KV heads.

```python
key_value_head_count: int
```

## depth

Channels per head.

```python
depth: int
```

## tokens

Total KV tokens across segments.

```python
tokens: int
```

## shape

Dimensions (key_value_head_count, depth, segments) read from the packed buffer header.

```python
shape: tuple[int, int, int]
```

## dtype

Data type of the packed KV-cache.

```python
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName
```

## nbytes

Size of the packed buffer in bytes.

```python
nbytes: int
```

## cdist

Compute distances between every row of `a` and every row of `b`.

```python
def cdist(
    a: _BufferType,
    b: _BufferType,
    /,
    metric: _MetricName = 'euclidean',
    *,
    threads: int = 1,
    dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType | None = None,
    out: _BufferType | None = None,
    out_dtype: _FloatTypeName | _ComplexTypeName | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> float | complex | Tensor | None
```

Similar to [`scipy.spatial.distance.cdist`](#cdist). Two vectors give a scalar. Matrices give a
[rows,columns] Tensor with a row per row of `a` and a column per row of `b`, or fill `out` and
return None. GPU tensors raise BufferError; pack `b` with dots_pack() and call the `_packed`
functions instead.

- `a`: First vector or matrix.
- `b`: Second vector or matrix, of the dtype of `a`.
- `metric`: Name of the distance, 'euclidean' by default.
- `threads`: Worker threads.
- `dtype`: Input dtype to read both buffers as, overriding the one they report.
- `out`: Output matrix for the distances.
- `out_dtype`: Result dtype, by default that of `out`, else float64, complex128 for complex input.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## dots_symmetric

Compute the Gram matrix, the symmetric all-pairs dot products, of the rows of [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors).

```python
def dots_symmetric(
    vectors: _BufferType,
    /,
    *,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    out: _BufferType | None = None,
    rows_begin: int | None = None,
    rows_end: int | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Only the upper triangle of the [count,count] result is guaranteed to be initialized. Returns a
new Tensor, or `out` when given. On a GPU the call returns before the kernel finishes.

- `vectors`: Input matrix of shape [count,depth], or a CUDA or ROCm Tensor from from_dlpack() to
run on that GPU.
- `dtype`: Input dtype to read [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors) as, overriding the one it reports.
- `out`: C-contiguous [count,count] output, required on the same GPU for GPU [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors).
- `rows_begin`: First row to compute, 0 when None.
- `rows_end`: One past the last row to compute, `count` when None; only the upper triangle
overlapping the rows in range is filled.
- `threads`: Worker threads, ignored on a GPU.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## hammings_symmetric

Compute the symmetric all-pairs Hamming-distance matrix of the rows of [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors).

```python
def hammings_symmetric(
    vectors: _BufferType,
    /,
    *,
    dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    out: _BufferType | None = None,
    rows_begin: int | None = None,
    rows_end: int | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Only the upper triangle of the [count,count] uint32 result is guaranteed to be initialized.
Returns a new Tensor, or `out` when given.

- `vectors`: Input matrix of shape [count,depth], with 'uint1' bits packed into uint8 bytes.
- `dtype`: Input dtype to read [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors) as, overriding the one it reports.
- `out`: C-contiguous [count,count] output of dtype uint32.
- `rows_begin`: First row to compute, 0 when None.
- `rows_end`: One past the last row to compute, `count` when None; only the upper triangle
overlapping the rows in range is filled.
- `threads`: Worker threads.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## jaccards_symmetric

Compute the symmetric all-pairs Jaccard-distance matrix of the rows of [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors).

```python
def jaccards_symmetric(
    vectors: _BufferType,
    /,
    *,
    dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    out: _BufferType | None = None,
    rows_begin: int | None = None,
    rows_end: int | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Only the upper triangle of the [count,count] result is guaranteed to be initialized. Returns a
new Tensor, or `out` when given.

- `vectors`: Input matrix of shape [count,depth], with 'uint1' bits packed into uint8 bytes.
- `dtype`: Input dtype to read [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors) as, overriding the one it reports.
- `out`: C-contiguous [count,count] output.
- `rows_begin`: First row to compute, 0 when None.
- `rows_end`: One past the last row to compute, `count` when None; only the upper triangle
overlapping the rows in range is filled.
- `threads`: Worker threads.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## angulars_symmetric

Compute the symmetric all-pairs angular-distance matrix of the rows of [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors).

```python
def angulars_symmetric(
    vectors: _BufferType,
    /,
    *,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    out: _BufferType | None = None,
    rows_begin: int | None = None,
    rows_end: int | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Only the upper triangle of the [count,count] result is guaranteed to be initialized. Returns a
new Tensor, or `out` when given. On a GPU the call returns before the kernel finishes.

- `vectors`: Input matrix of shape [count,depth], or a CUDA or ROCm Tensor from from_dlpack() to
run on that GPU.
- `dtype`: Input dtype to read [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors) as, overriding the one it reports.
- `out`: C-contiguous [count,count] output, required on the same GPU for GPU [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors).
- `rows_begin`: First row to compute, 0 when None.
- `rows_end`: One past the last row to compute, `count` when None; only the upper triangle
overlapping the rows in range is filled.
- `threads`: Worker threads, ignored on a GPU.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## euclideans_symmetric

Compute the symmetric all-pairs Euclidean-distance matrix of the rows of [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors).

```python
def euclideans_symmetric(
    vectors: _BufferType,
    /,
    *,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    out: _BufferType | None = None,
    rows_begin: int | None = None,
    rows_end: int | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Only the upper triangle of the [count,count] result is guaranteed to be initialized. Returns a
new Tensor, or `out` when given. On a GPU the call returns before the kernel finishes.

- `vectors`: Input matrix of shape [count,depth], or a CUDA or ROCm Tensor from from_dlpack() to
run on that GPU.
- `dtype`: Input dtype to read [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors) as, overriding the one it reports.
- `out`: C-contiguous [count,count] output, required on the same GPU for GPU [`vectors`](https://www.unum.cloud/docs/numkong/python/maxsim.md#vectors).
- `rows_begin`: First row to compute, 0 when None.
- `rows_end`: One past the last row to compute, `count` when None; only the upper triangle
overlapping the rows in range is filled.
- `threads`: Worker threads, ignored on a GPU.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## dots_pack

Pack a matrix with contiguous rows for repeated dot-product cross operations.

```python
def dots_pack(
    b: _BufferType,
    /,
    dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType | None = None,
    *,
    out: _BufferType | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> PackedMatrix
```

Packing fuses no transposition, so transpose [`w.T`](https://www.unum.cloud/docs/numkong/python/tensors-and-views.md#T) into a contiguous copy first. The pack serves
dots_packed(), angulars_packed(), euclideans_packed() and the `@` operator of Tensor.

- `b`: Source matrix of shape [columns,depth], or a CUDA or ROCm Tensor from from_dlpack() to pack
on that GPU.
- `dtype`: Packing dtype, inferred from `b` when None.
- `out`: C-contiguous Tensor on the GPU of `b` holding the pack, required for a GPU `b`, of at
least PackedMatrix.pack_size() bytes.
- `threads`: Worker threads, ignored on a GPU.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## dots_packed

Compute dot products between rows of `a` and a matrix packed by dots_pack(), like `a @ b.T`.

```python
def dots_packed(
    a: _BufferType,
    b: PackedMatrix,
    /,
    *,
    out: _BufferType | None = None,
    rows_begin: int | None = None,
    rows_end: int | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Returns a [rows,columns] Tensor, or `out` when given. On a GPU the call returns before the
kernel finishes.

- `a`: Query matrix of shape [rows,depth], on the device of `b`.
- `b`: Matrix packed by dots_pack(), of shape [columns,depth].
- `out`: C-contiguous [rows,columns] output, required on the same GPU for a GPU-packed `b`.
- `rows_begin`: First row of `a` to process, 0 when None.
- `rows_end`: One past the last row of `a` to process, every row when None.
- `threads`: Worker threads, ignored on a GPU.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## hammings_pack

Pack a matrix for repeated Hamming and Jaccard cross operations.

```python
def hammings_pack(
    b: _BufferType,
    /,
    dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType | None = None,
    *,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> PackedMatrix
```

The pack serves hammings_packed() and jaccards_packed().

- `b`: Source matrix of shape [columns,depth], with 'uint1' bits packed into uint8 bytes.
- `dtype`: Packing dtype, 'uint1' when None.
- `threads`: Worker threads.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## hammings_packed

Compute Hamming distances between rows of `a` and a matrix packed by hammings_pack().

```python
def hammings_packed(
    a: _BufferType,
    b: PackedMatrix,
    /,
    *,
    out: _BufferType | None = None,
    rows_begin: int | None = None,
    rows_end: int | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Returns a [rows,columns] uint32 Tensor, or `out` when given.

- `a`: Query matrix of shape [rows,depth].
- `b`: Matrix packed by hammings_pack(), of shape [columns,depth].
- `out`: C-contiguous [rows,columns] output of dtype uint32.
- `rows_begin`: First row of `a` to process, 0 when None.
- `rows_end`: One past the last row of `a` to process, every row when None.
- `threads`: Worker threads.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## jaccards_packed

Compute Jaccard distances between rows of `a` and a matrix packed by hammings_pack().

```python
def jaccards_packed(
    a: _BufferType,
    b: PackedMatrix,
    /,
    *,
    out: _BufferType | None = None,
    rows_begin: int | None = None,
    rows_end: int | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Returns a [rows,columns] Tensor, or `out` when given.

- `a`: Query matrix of shape [rows,depth].
- `b`: Matrix packed by hammings_pack(), of shape [columns,depth].
- `out`: C-contiguous [rows,columns] output.
- `rows_begin`: First row of `a` to process, 0 when None.
- `rows_end`: One past the last row of `a` to process, every row when None.
- `threads`: Worker threads.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## angulars_packed

Compute angular distances between rows of `a` and a matrix packed by dots_pack().

```python
def angulars_packed(
    a: _BufferType,
    b: PackedMatrix,
    /,
    *,
    out: _BufferType | None = None,
    rows_begin: int | None = None,
    rows_end: int | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Returns a [rows,columns] Tensor, or `out` when given. On a GPU the call returns before the
kernel finishes.

- `a`: Query matrix of shape [rows,depth], on the device of `b`.
- `b`: Matrix packed by dots_pack(), of shape [columns,depth].
- `out`: C-contiguous [rows,columns] output, required on the same GPU for a GPU-packed `b`.
- `rows_begin`: First row of `a` to process, 0 when None.
- `rows_end`: One past the last row of `a` to process, every row when None.
- `threads`: Worker threads, ignored on a GPU.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## euclideans_packed

Compute Euclidean distances between rows of `a` and a matrix packed by dots_pack().

```python
def euclideans_packed(
    a: _BufferType,
    b: PackedMatrix,
    /,
    *,
    out: _BufferType | None = None,
    rows_begin: int | None = None,
    rows_end: int | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Returns a [rows,columns] Tensor, or `out` when given. On a GPU the call returns before the
kernel finishes.

- `a`: Query matrix of shape [rows,depth], on the device of `b`.
- `b`: Matrix packed by dots_pack(), of shape [columns,depth].
- `out`: C-contiguous [rows,columns] output, required on the same GPU for a GPU-packed `b`.
- `rows_begin`: First row of `a` to process, 0 when None.
- `rows_end`: One past the last row of `a` to process, every row when None.
- `threads`: Worker threads, ignored on a GPU.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## attention_pack

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

```python
def attention_pack(
    k: _BufferType,
    v: _BufferType,
    /,
    key_offsets: _BufferType,
    key_lengths: _BufferType | None = None,
    depth: int | None = None,
    out: Tensor | None = None,
    threads: int = 1,
    *,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> AttentionPackedMatrix
```

Segment s keeps its `key_lengths[s]` keys, by default the whole slot, from row
`key_offsets[s]`.

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## attention_packed

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

```python
def attention_packed(
    q: _BufferType,
    kv: AttentionPackedMatrix,
    /,
    query_offsets: _BufferType,
    out: Tensor | None = None,
    scale: float | None = None,
    keys_before: int | None = None,
    keys_after: int | None = None,
    log_sum_exp: _BufferType | None = None,
    threads: int = 1,
    *,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

Query row r of q queries over k keys sees the keys from `r + k - q - keys_before` through
`r + k - q + keys_after`, each side unbounded at None.

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## attention_packed_gradients

Compute the bf16 query, key and value gradients of [`attention_packed`](#attention_packed) under the same band.

```python
def attention_packed_gradients(
    q: _BufferType,
    kv: AttentionPackedMatrix,
    /,
    query_offsets: _BufferType,
    output: _BufferType,
    output_gradient: _BufferType,
    log_sum_exp: _BufferType,
    scale: float | None = None,
    keys_before: int | None = None,
    keys_after: int | None = None,
    query_gradient: Tensor | None = None,
    threads: int = 1,
    *,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> tuple[Tensor, Tensor, Tensor]
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## attention_rope

Rotate every channel pair of each head by the per-token angle grids, in place unless `out` is given.

```python
def attention_rope(
    x: _BufferType,
    cos: _BufferType,
    sin: _BufferType,
    /,
    head_count: int,
    depth: int,
    *,
    out: _BufferType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> None
```

- `x`: `[rows,head_count*depth]`, float32, bfloat16, or e4m3.
- `cos`: `[rows,depth/2]` float32 angle grid, shared across heads.
- `sin`: `[rows,depth/2]` float32 angle grid, shared across heads.
- `head_count`: Number of heads per token.
- `depth`: Even number of channels per head.
- `out`: Output, same shape and dtype as `x`, may alias `x`, defaulting to `x`.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.