Primitives

StableTested and in active use. Breaking changes will be called out.

Around eighty leaf modules live in noeira.nn.primitives. Every one implements Module with explicit forward and vjp, and most have both a CPU and a GPU kernel.

Import paths follow the file name:

from noeira.nn.primitives.linear import Linear
from noeira.nn.primitives.linear_relu import LinearReLU
from noeira.nn.primitives.conv2d import Conv2D
from noeira.nn.primitives.layer_norm import LayerNorm

Reference

• Layer NormalizationBa, Kiros, Hinton · 2016 · arXiv:1607.06450

Linear and fused variants

Module	Notes
Linear[IN, OUT]	dense layer with bias
LinearReLU[IN, OUT]	fused — one GPU kernel
LinearTanh · LinearSigmoid · LinearSwish · LinearMish	other fused activations
LinearAct	generic fused-activation form
BlockLinear	block-diagonal, for ensembles
TiedLinear	weights tied to another layer — used for GPT output projection
NoisyLinear	learned parameter noise; replaces ε-greedy in Noisy DQN
ZeroLinear	zero-initialized output layer

The fused variants are what the GPU examples use throughout. Linear followed by ReLU is two kernels; LinearReLU is one, and in an inner training loop evaluated millions of times that difference is real.

FeedForwardGELU[S, D, FF] is a transformer FFN, Tok[Linear] → GELUTanh → Tok[Linear], as one module. On NVIDIA it computes fc1’s GELU in the GEMM’s epilogue, so the separate GELU pass over the hidden activations is gone. The GPT and the classifier ViT use it; elsewhere it runs its three children as the Sequential would.

Where the GEMMs run on NVIDIA

On NVIDIA, Linear, LinearAct and Conv2D choose their GEMM per shape, by a rule measured on the shapes noeira’s agents and examples actually use:

• Backward always runs on cuBLAS (noeira/nn/core/cublas_gemm.mojo), on the unpadded operands. MAX’s GEMM rejects transpose_a and falls back to an allocating path unless N % 128 == 0, so the old backward transposed and padded. An 8 → 64 RL trunk did 32× its work.
• Forward runs on cuBLAS too, except where MAX measured faster: batch 1 on an aligned shape, very large batches with a mid-sized K, or a 2²⁵-element weight at a small batch.
• Epilogues. The bias, and for LinearReLU the bias and the ReLU, go in a cuBLASLt epilogue (noeira/nn/core/cublaslt_gemm.mojo) on the shapes where that measured faster. The bf16 path fuses the bias, and FeedForwardGELU’s GELU, the same way.
• Precision. Full fp32 below 2²⁴ multiply-adds per layer, TF32 above. With every GEMM in TF32, PPO stalled on Hopper in 7 of 20 seeds, against 2 of 20 before.

On an RTX 5090, fwd+bwd was faster than the old path on all 119 measured shapes (median 2.9×). End to end: PPO LunarLander 19.8k → 25.4k env steps/s; SAC HalfCheetah (32 envs, 200k steps) 200 → 87 s. Apple keeps MAX’s GEMMs.

For A/B runs, -D NN_GEMM_PATH=auto|max|cublas|fwdmax|bwdmax forces a route, and -D NN_CUBLAS_FP32 turns TF32 off.

Weight tying is the dominant lever in small language models

On character-level GPT, tying the output projection to the embedding was worth more than any other single change. TiedLinear exists for exactly this.

Convolution

Module	Notes
Conv2D	2D convolution, GEMM-backed via im2col
Conv2DTranspose	transposed convolution — decoders (DreamerV3 pixel)
MaxPool2D · AvgPool2D	pooling
Upsample2x	nearest-neighbour upsampling
ConvRMSNorm	RMS normalization for conv activations
ToNCHW · Flatten · Transpose2D	layout helpers

Unlike Linear, Conv2D has no fused-activation form: follow it with an activation module.

On NVIDIA, Conv2D picks its path per shape. Its three im2col GEMMs (forward, dW, d_col) run unpadded on cuBLAS. Large stride-1 convolutions, OC = 1 and small 1×1 heads go to cuDNN (noeira/nn/core/cudnn_conv.mojo). Conv2D.use_cudnn holds that rule, measured on 80 shapes from every agent and example: none got slower, median 1.49×. cuDNN’s algorithm is the top of its own heuristic ranking among its deterministic GEMM algorithms. Timing-based selection (cudnnFind) picked different algorithms in different processes, so one binary trained differently from run to run.

On an RTX 5090, ResNet-20 went from 6.16 to 2.84 ms per step, level with torch.compile plus CUDA graphs. The CIFAR CNN went from 2.36 to 1.35 ms (PyTorch: 1.33).

Attention and transformers

Module	Notes
ScaledDotProductAttention	self-attention over one packed [Q|K|V] input
ScaledDotProductAttentionQKV	the same, reading the QKV projection’s token-major output directly — no QKVToMajor pass
MaskedAttention	causal masking — one [SEQ, SEQ] bias shared by the batch
CrossAttention[DIM, N_HEADS, Q_LEN, KV_LEN, MASKED]	separate query and key/value streams, optional per-sample key padding mask
SelfAttentionPos · SelfAttentionPosMasked	CrossAttention at Q_LEN == KV_LEN — DETR self-attention
RoPE[SEQ, N_HEADS, HEAD_DIM, THETA]	rotary position embedding, applied to q and k
RepeatKVHeads[SEQ, N_KV, REP, HEAD_DIM]	broadcasts K/V heads for grouped-query attention
QKVToMajor	QKV layout transform
DecoderBlock	pre-norm transformer decoder block
ConditionalTransformerBlock	conditioned variant
Embedding · SinusoidalPosAdd · SinusoidalPos1DTokens · SinusoidalPos2DTokens	token and position embeddings
LearnedQueries[IGNORE_DIM, N, D] · LearnedTokens · BroadcastTokens	latent-query patterns
PixelShuffle[GRID, EMBED, SCALE]	folds a patch grid’s space into channels
SwiGLU[HIDDEN] · Gate · Modulate	gating and modulation

The modules the imitation policies (ACT, SmolVLA) are built from:

CrossAttention takes q, k and v as three inputs (a fourth, the key padding mask, when MASKED=True). This is what DETR-style models need and ScaledDotProductAttention cannot express: queries attending to a memory of a different length, and a positional embedding added to q and k but not to v. The mask polarity is 1.0 = attend, 0.0 = ignore — the inverse of PyTorch’s key_padding_mask. A fully masked row produces a zero context vector rather than NaN.

RoPE implements the split-halves pairing (channel j with j + HEAD_DIM/2) that Hugging Face safetensors checkpoints use, not GGUF’s interleaved pairs. The wrong pairing loads and runs and computes a different function; tests/nn/test_rope_vs_max_kernel.mojo checks it against MAX’s own RoPE kernel.

RepeatKVHeads maps output head g to K/V head g // REP, matching Hugging Face’s repeat_kv (SmolLM2 has 15 query heads over 5 K/V heads). Its backward sums the REP copies’ gradients.

PixelShuffle turns [B, GRID², EMBED] into [B, (GRID/SCALE)², EMBED·SCALE²] — SmolVLM’s connector, 1024 tokens × 768 into 64 tokens × 12288. It is a pure permutation, so its gate is the adjoint identity.

SwiGLU is the parameter-free middle of a gated FFN: Sequential[Linear[d, 2*hidden], SwiGLU[hidden], Linear[hidden, d]].

LearnedQueries is a learned constant token sequence whose input is read only for its batch size — the DETR decoder’s queries. It draws its initial values from N(0, 1) (EMB_STD = 1.0), as nn.Embedding does; a narrower draw left ACT’s decoder returning the same action at every position of the chunk.

QKV layout mismatches scramble attention silently

A wrong QKV layout produces attention that runs, trains, and is wrong. Worse, a numerical parity test against a reference will not catch it if the test does not exercise causality — the shapes match and the values look plausible. Test masked attention with a sequence where position t must not see t+1. The same holds for RoPE’s pairing, RepeatKVHeads’ grouping and PixelShuffle’s permutation: every wrong choice is shape-legal.

Attention performance

Self-attention is fused on the GPU by default. ScaledDotProductAttention and ScaledDotProductAttentionQKV take FLASH=True. When there is no attention dropout and the head dim is at most 64 and a multiple of the tile width (32 on NVIDIA, 16 on Apple), they run FlashAttention-2 forward and backward. The forward keeps only the output and each row’s log-sum-exp, with no [B, H, S, S] score matrix. The backward recomputes the probabilities and uses no atomics, so an eager run and a CUDA-graph replay stay bit-identical. Which kernels run depends on the activation dtype and the device:

Activations	NVIDIA	Apple
fp32	tensor cores in TF32 (flash_attention_tf32.mojo)	fp32 kernels (flash_attention.mojo)
bf16	tensor cores, bf16 operands, fp32 accumulators (flash_attention_mma.mojo)	—

-D NN_ATTN_PATH=no_tf32 keeps NVIDIA on the strict fp32 kernels. The TF32 kernels round their operands to TF32 rather than truncating them, as the tensor core would, which would bias every softmax weight down. Dropout, wider heads and the CPU keep the pack / bmm / softmax / bmm path.

On an RTX 5090, the TinyShakespeare GPT (6 × 384, sequence 256, batch 64) trains a step in 19.4 ms in fp32 and 10.74 ms in bf16. PyTorch 2.12 takes 21.35 and 10.81 ms (reduce-overhead, same box). MAX’s own FA2 kernel was read first: it admits head dim 64 only on sm90 / sm100 / A100, by GPU name, and it has no backward.

CrossAttention has a second, inference-only GPU forward: a fused online-softmax kernel that never materialises the score matrix. It is off by default and selected with set_attr["fused_attention"](1.0), which every container forwards. It writes no softmax weights, so vjp after a fused forward raises. The SmolVLA deploy turns it on for the frozen vision tower (--no-fused-vision turns it off in examples/so101/smolvla_so101_deploy_real.mojo).

On NVIDIA, at head dim 64 and without a mask, the fused flag routes to MAX’s own FA2 flash-attention kernel instead (xa_fused_routes_to_max in noeira/nn/primitives/cross_attention.mojo). Measured on a Jetson Orin at SigLIP’s shape (batch 1, 12 heads, 1024 × 1024 tokens, head dim 64):

Path	ms per layer
two-pass (scores materialised)	8.42
noeira’s fused kernel	4.13
MAX FA2 kernel, routed	1.26

The routed path is TF32

MAX’s kernel runs both matmuls on the tensor cores in TF32: about 3.7e-3 standard units against a float64 reference at that shape, where the fp32 kernels sit near 7e-6. That is the band the network’s GEMMs already run in on CUDA, and tests/nn/test_cross_attention_gpu_shapes.mojo carries a TF32 band for exactly the routed shapes. Metal, ACT’s head dim 32 and masked attention keep noeira’s own kernel.

The benchmark behind the table is benchmarks/cross_attention_bench.mojo (pixi run -e jetson cross-attn-bench-jetson builds it for the Orin).

Recurrence

Module	Notes
LSTMCell	single LSTM step
LSTMSeq	unrolled sequence
GRUCell	GRU step — the RSSM recurrence in DreamerV3

Do not spectral-clamp a tanh recurrence

Constraining the spectral norm of a tanh recurrent weight matrix damages learning rather than stabilising it. The saturating nonlinearity already bounds the state.

Normalization

Module	Notes
LayerNorm · LayerNormNoAffine	the deep-RL default
RMSNorm[DIM, EPS] · ConvRMSNorm	cheaper, no mean subtraction
BatchNorm1D · BatchNorm2D	running statistics
SimNorm	simplicial normalization — TD-MPC2 latents
MinMaxNorm	bounded normalization

RMSNorm’s default epsilon is 1e-4. SmolLM2 and Llama use 1e-5, and the difference shows in a parity check while passing every shape and NaN test — pass RMSNorm[DIM, EPS=…] when porting a model.

Prefer LayerNorm in deep RL

Value scales drift over training in ways batch statistics cannot track, and RL minibatches are not i.i.d. LayerNorm is the safer default — it fixed critic instability on Humanoid (SAC) and runaway Q-values in MBPO.

If you do use BatchNorm, remember its running statistics are state, not parameters. A hard copy to a target network that misses them yields a network that evaluates catastrophically while looking fine.

Activations

ReLU, Tanh, Sigmoid, GELU, GELUTanh, ELU, Mish, SiLU, Swish, Symlog, Clamp, Scale.

from noeira.nn.primitives.activations import ReLU, Tanh, Sigmoid

GELU is the erf form

GELU is PyTorch’s exact GELU (nn.GELU(), Hugging Face’s "gelu"). GELUTanh is the tanh approximation: JAX’s default, Hugging Face’s "gelu_pytorch_tanh". When porting a model, pick the one its config names. The two differ by up to about 5e-4. A parity check at a loose band passes either one, and so does every shape and NaN test.

ELU (alpha 1) is PyTorch’s nn.ELU, the activation of the rsl_rl-style locomotion networks.

Symlog — sign(x)·log(1+|x|) — compresses unbounded targets and appears throughout DreamerV3.

RL-specific operations

These are why nn is not a general-purpose framework wearing an RL hat:

Module	Used by
RSample	reparameterised Gaussian sampling — SAC, PPO
DuelingHead · DuelingHeadC51	Dueling DQN, Rainbow
GatherActionSlice	select the taken action’s Q-value
MSEPerSample	per-sample loss for prioritized replay weights
StopGrad · StopGradParams	detach in a loss graph
SIGReg	SIGReg regularizer — the LeWM world-model work
TokenMean · Slice · Concat · Reduce	plumbing

Utility

Add, BiasAdd, Elementwise, BinaryElementwise, Dropout, HashDropout, GradGate, GatherCols, ReduceMax, Concat, Slice, SpaceTimeTranspose, MAEReplacer.

Dropout draws its masks from a Philox stream keyed by a compile-time seed. Inside a Repeat, copy i draws from subsequence i. Before this was fixed, the N copies shared one seed and every GPT block dropped the same entries as the others. set_attr["rng_stream"](k) moves a replica onto stream k; the distributed trainer gives rank r stream r. RepeatConditional does not give its copies separate subsequences; use HashDropout there. Its mask is a hash of its own seed, a counter and the element index, so the backward redraws the mask instead of caching it, and CPU and GPU draw the same masks.

GradGate is the identity going forward and a stop-gradient you can switch at run time (set_attr["stop_grad"](1)). Switched off, it trains bit-identically to a graph without it.

A note on GEMM

transpose_b GEMM miscomputes at N = 1 on GPU

A GPU matrix multiply with transpose_b and N = 1 silently produces wrong results. It was found via Conv2D forward with a single output channel.

If you add a call site where the second dimension can be 1, either guard it or ensure N ≥ 2. The failure is silent.

Next

• Combinators and models — composing these
• Modules and Params — writing your own leaf
• Optimizers and losses — what trains them