Linear algebra & NN

10 functions — 10 documented.

matmul¶

Matrix product of two tensors.

The behaviour depends on the dimensionality of the arguments:

• If both are 2-D, the ordinary matrix product.
• If either is more than 2-D, the leading axes are batch and broadcast against each other; the last two are multiplied.
• The inner dimensions must agree: (…, n, k) against (…, k, m) gives (…, n, m).

Six pipelines sit behind this one operator — gemv, gemm_naive, gemm_tiled, gemm_db, gemm_reg and gemm_split_k_reduce — and it is the only operator in vkML that creates more than one. Which runs is decided per dispatch, in this order:

• GEMV, when explicitly selected. One workgroup per output element. The tiled kernel collapses to ceil(M/32) workgroups when N=1, which on the development GPU is 128 against 288 concurrent slots — 44% occupancy, where throughput falls off a cliff. GEMV restores the grid to M·N.
• Naive, tiled or register-blocked, from VKML_GEMM_KERNEL; the register-blocked kernel is the default.
• A forced fall back to naive when the device cannot run the blocked kernel. Both blocked kernels hardcode 256 invocations — deliberately, so the three variants stay comparable — and Vulkan guarantees only 128.
• Split-K, for the register-blocked kernel only, when the K dimension is long relative to the output.

The register block is BM=32, BN=32, RM=2, RN=2 with BK=32. BK is 32 rather than 16 for two measured reasons: it halves the K-tile count, which costs one fewer carry-stack level and therefore fewer registers; and it makes each block a 32-element sequential sum, exactly matching kPairwiseBlock in src/backend/cpu/reduce.h, so the two backends fold K with the same structure.

Parameters

input (Tensor)

The left operand.

other (Tensor)

The right operand. Its second-to-last axis must match input's last axis.

Returns

The product, with batch axes broadcast.

Example

copy>>> a = vkml.tensor(np.random.rand(64, 128).astype(np.float32))
>>> b = vkml.tensor(np.random.rand(128, 32).astype(np.float32))
>>> vkml.matmul(a, b).shape
(64, 32)

From the header

Matrix multiply, following torch.matmul for ranks 1-4: leading axes are treated as batch and broadcast against each other.

Implementation

See also conv2d, im2col

softmax¶

Apply the softmax function along one axis.

Computes exp(xᵢ − max(x)) / Σ exp(xⱼ − max(x)) along dim.

One workgroup per row, three passes inside it: the maximum over the axis, the sum of exp(x − max), then the normalised write. All three share a single workgroup, so the two reductions synchronise with barrier() rather than needing separate dispatches and a global barrier — which is the whole reason softmax is one kernel and not three. The shader carries seven barriers, more than any other in the project.

The maximum is subtracted before exponentiating so a large logit cannot overflow: exp reaches +inf above about 88.72 in float32.

Parameters

input (Tensor)

Any shape, float dtype.

dim (int = -1) optional

Axis to normalise over. Negative counts from the end.

Returns

A tensor of the same shape whose values along dim sum to 1.

Example

copy>>> x = vkml.tensor(np.array([[1.0, 2.0, 3.0]], dtype=np.float32))
>>> vkml.softmax(x, -1).numpy()
array([[0.09003057, 0.24472848, 0.66524094]], dtype=float32)

Implementation

See also log_softmax, cross_entropy, sum

log_softmax¶

The logarithm of softmax, computed without forming the softmax.

x − max(x) − log(Σ exp(x − max(x))), sharing the softmax kernel and its three-pass structure — the difference is only what the final pass writes.

Computing it this way is not an optimisation. log(softmax(x)) underflows: a probability reaches exactly 0 in float32 once the logit gap passes about 90, and log(0) is -inf, which then poisons every gradient in the batch.

Parameters

input (Tensor)

Any shape, float dtype.

dim (int = -1) optional

Axis to normalise over.

Returns

A tensor of the same shape.

Example

copy>>> x = vkml.tensor(np.array([[1.0, 2.0, 3.0]], dtype=np.float32))
>>> vkml.log_softmax(x, -1).numpy()
array([[-2.407606  , -1.4076059 , -0.40760595]], dtype=float32)

Implementation

See also softmax, cross_entropy, kl_div

layer_norm¶

Normalise over the trailing axes to zero mean and unit variance.

Takes a count of trailing axes rather than a shape: normalized_axes=1 normalises over the last axis, 2 over the last two. No weight or bias — vkML's layer_norm is the normalisation alone, and vkml.nn.LayerNorm applies the affine transform on top of it.

Parameters

input (Tensor)

Any shape, float dtype.

normalized_axes (int = 1) optional

How many trailing axes to normalise over.

eps (float = 1e-5) optional

Added to the variance before the square root.

Returns

A tensor of the same shape and dtype.

Example

copy>>> x = vkml.tensor(np.array([[1.0, 2.0, 3.0]], dtype=np.float32))
>>> vkml.layer_norm(x, 1, 1e-5).numpy()
array([[-1.2247356,  0.       ,  1.2247356]], dtype=float32)

From the header

Standardises over the last normalized_axes axes: subtract the mean, divide by the standard deviation. No affine term -- nn.LayerNorm applies its own weight and bias with mul/add, which the autograd handles without a special case.

Composed from mean/sub/square/rsqrt rather than fused. That is the two-pass algorithm, which is the numerically sound one: computing the variance as E[x^2] - E[x]^2 in one pass cancels catastrophically when the mean is large relative to the spread. A fused kernel would save bandwidth, not accuracy, and is deferred until a profile says it matters.

normalized_axes counts trailing axes rather than naming a shape, which is a deliberate divergence from torch's normalized_shape. It carries the same information for every legal call and needs no shape validation.

Implementation

See also rms_norm, batch_norm, rsqrt

rms_norm¶

Normalise over the trailing axes by root-mean-square, without centring.

x / sqrt(mean(x²) + eps). Unlike layer_norm it does not subtract the mean, which makes it cheaper — one reduction instead of two — and is why transformer implementations increasingly prefer it.

Parameters

input (Tensor)

Any shape, float dtype.

normalized_axes (int = 1) optional

How many trailing axes to normalise over.

eps (float = 1e-5) optional

Added to the mean square before the root.

Returns

A tensor of the same shape and dtype.

Example

copy>>> x = vkml.tensor(np.array([[1.0, 2.0, 3.0]], dtype=np.float32))
>>> vkml.rms_norm(x, 1, 1e-5).numpy()
array([[0.46290955, 0.9258191 , 1.3887286 ]], dtype=float32)

From the header

As layer_norm but without centring: divides by the root mean square. Used by Llama-family models, where dropping the mean subtraction is most of the speedup and costs nothing measurable in quality.

DEFAULT eps DIVERGES FROM TORCH, deliberately. torch.nn.functional.rms_norm defaults eps to finfo(dtype).eps (1.19e-7 for f32), which is small enough to be a no-op except for an exactly-zero input. vkml uses 1e-5, matching layer_norm here and what Llama-family implementations actually ship. The two are otherwise identical: pass eps explicitly and the results agree to 7e-7, which is how the parity tests compare them.

Implementation

See also layer_norm, rsqrt

batch_norm¶

Normalise each channel by supplied statistics, then scale and shift.

(x − mean) / sqrt(variance + eps) · weight + bias, with the statistics passed in rather than computed. That is deliberate: this operator does not decide between batch statistics and running statistics, so it has no train/eval mode and no hidden state. vkml.nn.BatchNorm2d owns that decision and the running buffers, and calls this.

Parameters

input (Tensor)

(N, C, …) — channels on axis 1.

mean (Tensor)

Per-channel mean, (C,).

variance (Tensor)

Per-channel variance, (C,).

weight (Tensor = undefined) optional

Per-channel scale, (C,). Omit for none.

bias (Tensor = undefined) optional

Per-channel shift, (C,). Omit for none.

eps (float = 1e-5) optional

Added to the variance before the root.

Returns

A tensor of the same shape as input.

From the header

Applies a batch normalisation given the statistics to use: (input - mean) / sqrt(variance + eps) * weight + bias.

mean and variance are rank-1 with one entry per channel, where the channel is axis 1. weight and bias are optional and likewise per-channel.

TAKES THE STATISTICS; IT DOES NOT CHOOSE OR UPDATE THEM. Whether to use the batch's own statistics or the running estimate is a property of training mode, and updating the running estimate is a mutation across calls -- both belong to the nn module, which owns that state, exactly as the optimisers own theirs. Keeping this function pure is what makes it directly comparable against torch.nn.functional.batch_norm.

A CALLER COMPUTING BATCH STATISTICS MUST USE THE BIASED VARIANCE (divide by N, not N-1) here, while updating a running estimate with the UNBIASED one. That is torch's behaviour, verified, and the asymmetry is deliberate on their part: the biased estimate is the right normaliser for the batch in hand, and the unbiased one the right estimator of the population. Using one for both makes evaluation drift away from training as the running estimate converges to the wrong value -- invisible in a single-step comparison.

Implementation

See also layer_norm, rms_norm

dropout¶

Randomly zero elements with probability p, scaling the rest.

Inverted dropout: survivors are scaled by 1/(1−p) during training so the expected value is unchanged and inference needs no compensation.

The mask comes from the same counter-based Philox generator as rand, which is why the signature takes (seed, offset) rather than reading hidden state. The mask is a pure function of those, so it is reproducible across runs and identical on both backends — and you must advance offset between steps or every step drops the same elements.

Parameters

input (Tensor)

Any shape, float dtype.

p (float)

Probability of zeroing an element. 0.0 is the identity.

seed (int)

Identifies the random stream.

offset (int = 0) optional

Position in the stream. Advance between steps.

training (bool = True) optional

When False, returns the input unchanged.

Returns

A tensor of the same shape and dtype.

From the header

Zeroes each element independently with probability p and scales the rest by 1 / (1 - p), following torch.nn.functional.dropout.

The scaling happens at training time -- "inverted dropout" -- so that evaluation is the identity and needs no compensating factor. training = false returns the input unchanged, which is why the flag is a parameter rather than something the caller branches on.

The mask comes from rand(seed, offset), so the caller owns reproducibility: the same seed and offset give the same mask. A training loop must advance the offset between steps, or every step drops the same elements.

Implementation

See also rand

conv2d¶

Apply a 2-D convolution over a batch of images.

Lowered to im2col followed by matmul, not implemented as a direct convolution. src/api/ops.cpp builds it in four steps:

• im2col expands (N, C, H, W) into (N, C·kh·kw, L), where L is the number of window positions.
• The weight is reshaped from (C_out, C, kh, kw) to (C_out, C·kh·kw).
• matmul broadcasts the batch axis, so the weights are shared across the batch without a copy.
• The bias is reshaped to (C_out, 1, 1) so it broadcasts across batch and space with stride 0 rather than being materialised.

The consequence is that every GEMM improvement reaches convolution for free — and so does every GEMM limitation.

Parameters

input (Tensor)

(N, C_in, H, W).

weight (Tensor)

(C_out, C_in, kh, kw).

bias (Tensor = undefined) optional

(C_out,). Omit for none.

stride (Sequence[int] = [1, 1]) optional

Step in (H, W).

padding (Sequence[int] = [0, 0]) optional

Zero padding on both sides of (H, W).

dilation (Sequence[int] = [1, 1]) optional

Spacing between kernel elements.

Returns

(N, C_out, H_out, W_out).

Example

copy>>> x = vkml.tensor(np.random.rand(8, 3, 32, 32).astype(np.float32))
>>> w = vkml.tensor(np.random.rand(16, 3, 3, 3).astype(np.float32))
>>> vkml.conv2d(x, w, stride=[1, 1], padding=[1, 1]).shape
(8, 16, 32, 32)

From the header

2D convolution, following torch.nn.functional.conv2d.

input is (N, C_in, H, W), weight is (C_out, C_in, kernel_h, kernel_w), and bias is (C_out) or undefined.

Implemented as convolution-as-GEMM: im2col to lay each window out as a column, then one matmul against the flattened weights. That reuses the tuned matmul instead of needing a direct convolution kernel, and its gradient falls out of the gradients for im2col, matmul and reshape -- there is no conv-specific backward rule anywhere.

GROUPS ARE NOT SUPPORTED. Grouped and depthwise convolution need either a batched matmul over the group axis or a separate kernel, and nothing in the current roadmap uses them. Passing a weight whose input-channel extent disagrees with the input's is rejected rather than silently misinterpreted.

Implementation

See also im2col, col2im, matmul, max_pool2d

max_pool2d¶

Maximum over each sliding window, per channel.

Accepts a non-contiguous input. The kernel addresses planes through the same Operand stride machinery every other kernel uses, rather than assuming a packed layout — so a transposed or sliced tensor pools correctly without a materialising copy first. A CONTIGUOUS specialisation constant lets the packed case skip the stride arithmetic entirely.

Parameters

input (Tensor)

(N, C, H, W).

kernel (Sequence[int])

Window size in (H, W).

stride (Sequence[int] = [0, 0]) optional

Step. Defaults to the kernel size.

padding (Sequence[int] = [0, 0]) optional

Padding on both sides of (H, W).

dilation (Sequence[int] = [1, 1]) optional

Spacing between window elements.

Returns

(N, C, H_out, W_out).

Example

copy>>> x = vkml.tensor(np.arange(16, dtype=np.float32).reshape(1, 1, 4, 4))
>>> vkml.max_pool2d(x, [2, 2], [2, 2]).numpy()
array([[[[ 5.,  7.],
         [13., 15.]]]], dtype=float32)

From the header

2D max pooling, following torch.nn.functional.max_pool2d.

PADS WITH -INFINITY, not zero. A padded window must not report 0 as its maximum when every real element is negative, which is why this needs its own kernel rather than composing as max over im2col -- im2col pads with zero (verified against torch; avg_pool2d composes precisely because zero padding is what it wants).

Its gradient goes to exactly one position per window -- the first maximum in row-major order within the window, matching torch -- not split among ties.

Implementation

See also avg_pool2d, conv2d, amax

avg_pool2d¶

Mean over each sliding window, per channel.

Parameters

input (Tensor)

(N, C, H, W).

kernel (Sequence[int])

Window size in (H, W).

stride (Sequence[int] = [0, 0]) optional

Step. Defaults to the kernel size.

padding (Sequence[int] = [0, 0]) optional

Padding on both sides of (H, W).

Returns

(N, C, H_out, W_out).

Example

copy>>> x = vkml.tensor(np.arange(16, dtype=np.float32).reshape(1, 1, 4, 4))
>>> vkml.avg_pool2d(x, [2, 2], [2, 2]).numpy()
array([[[[ 2.5,  4.5],
         [10.5, 12.5]]]], dtype=float32)

From the header

2D average pooling, following torch.nn.functional.avg_pool2d with count_include_pad=True: padded positions contribute zero and are counted in the divisor.

Composed from im2col and a mean, because zero padding is exactly the semantics wanted here. Dilation is not accepted -- torch's avg_pool2d has no such parameter, and offering one would invent a behaviour to match.

Implementation

See also max_pool2d, mean, conv2d