Attention with bounded-memory control: kv memory compressed into a fixed set of slots.
linear attentionsequence mixer
Residual add fused with RMS norm; returns the normalized sum and the new residual stream.
fusednormalization
Scaled matrix product accumulated onto a scaled input.
feature mixerlinear
Multi-head attention with grouped kv heads (q_heads must divide evenly).
attentionsequence mixer
Attention over packed variable-length sequences, each attended independently.
attentionsequence mixervarlen
Based linear attention: a second-order Taylor approximation of the softmax kernel.
linear attentionsequence mixer
The classic fused bias activation: tanh-approximate GELU of a biased input.
activationfused
BitLinear: RMS norm, int8 activation quant, ternary weight quant, then the projection.
fusedlinearnormalizationquantized
chunk_global_cumsum179 cases Inclusive sequence-wide cumsum in either FLA layout, optionally reversed and scaled.
reduction
chunk_local_cumsum193 cases Inclusive cumsum reset at chunk boundaries, in either FLA layout.
reduction
A gated delta rule reading corrections through an auxiliary key p while writing through k.
linear attentionsequence mixer
Softmax cross-entropy with ignore-index and label smoothing, matching torch semantics.
loss
DeltaNet: a linear-attention state updated by the error-correcting delta rule.
linear attentionsequence mixer
Values corrected by a strictly causal delta rule in attention space, then attended normally.
attentionsequence mixer
Dynamic Tanh, an elementwise normalization replacement.
activationnormalization
Embedding lookup with gradients accumulated into the table.
linear
Causal softmax attention with a learned per-head log forget gate on the scores.
attentionsequence mixer
gated_delta_product203 cases Gated DeltaProduct: several delta-rule writes per step under one per-head forget gate.
linear attentionsequence mixer
gated_delta_rule783 cases The delta rule with a per-head log forget gate applied to the state each step.
linear attentionsequence mixer
Gated Oja's rule: per-value-channel decay with Oja's normalizing term keeping memory bounded.
linear attentionsequence mixer
GEGLU gating with the tanh-approximate GELU.
activationfeature mixerfused
Gaussian Error Linear Unit, exact or tanh-approximate.
activation
Gated linear attention with a per-key-dimension log forget gate.
linear attentionsequence mixer
Group normalization: standardize within each of num_groups channel groups, then scale and shift.
normalization
group_norm_linear1,848 cases Group norm fused with a linear projection.
feature mixerfusedlinearnormalization
Per-token on-policy GRPO loss with a KL penalty against the reference policy.
fusedloss
Clipped off-policy GRPO loss with a KL penalty against the reference policy.
fusedloss
Gated slot attention: two chained gated linear-attention passes through softmaxed slots.
linear attentionsequence mixer
hadamard_transform121 cases Sylvester Hadamard transform along the last dimension (a power of two).
linear
The HGRN elementwise gated recurrence.
linear attentionsequence mixer
int8_int2_matmul254 cases Integer matrix product with int8 activations and ternary int2 weights.
linearquantized
Identity-plus-low-rank delta rule with explicit recurrent state.
linear attentionsequence mixer
Generalized Jensen-Shannon divergence between student and constant teacher log-probs.
loss
Kimi Delta Attention: the delta rule with a per-key-dimension log forget gate.
linear attentionsequence mixer
KDA's per-dimension log forget gate transform.
activation
kda_gate_cumsum1,078 cases The KDA gate transform followed by a chunk-local inclusive cumsum.
activationfusedreduction
KL divergence from log-space inputs to a constant target distribution.
loss
L2 normalization along the last dim.
normalization
Layer normalization: standardize the last dim, then scale and shift.
normalization
layer_norm_gated716 cases Layer norm scaled by a swish or sigmoid gate.
activationfusednormalization
layer_norm_linear1,723 cases Layer norm fused with a linear projection.
feature mixerfusedlinearnormalization
layer_norm_linear_quant1,720 cases BitNet-style linear on a layer-normed input: int8/ternary fake quant around the matmul.
fusedlinearnormalizationquantized
layer_norm_swish_linear1,723 cases Layer norm, swish gate, then a linear projection, in one fused op.
activationfeature mixerfusedlinearnormalization
Lightning attention: linear attention with a fixed per-head, per-layer decay slope.
linear attentionsequence mixer
Causal linear attention, optionally normalized by the running key sum.
linear attentionsequence mixer
linear_cross_entropy527 cases Cross entropy fused with the lm-head projection, never materializing the logits.
fusedlinearloss
Generalized JSD between student and teacher lm-head outputs, fused with both projections.
fusedlinearloss
KL divergence between student and constant target lm-head outputs, fused with both projections.
fusedlinearloss
Paired Llama 4 rotary embedding, frequencies as real/imaginary pairs instead of complex.
positional
Gated linear attention through a hierarchical mask, one learned scale per log2 level.
linear attentionsequence mixer
Temperature-scaled log-sigmoid.
activation
Log-sum-exp over the last dimension, optionally after scaling.
reduction
Matrix product with an optional fused activation.
feature mixerlinear
Mean-pool consecutive, non-overlapping sequence chunks.
reduction
MesaNet: each step solves a ridge regression against gated key/value statistics.
linear attentionsequence mixer
One MesaNet decode step with explicit recurrent statistics and a conjugate-gradient solve.
linear attentionsequence mixer
Mixture of experts over SwiGLU MLPs, routed to the top-k experts by softmax probability.
activationfeature mixerfusedlinear
multi_token_attention421 cases Causal softmax over raw scores, a conv2d mixing attention maps, then re-masking the future.
attentionfusedsequence mixer
neighborhood_attn574 cases Attention restricted to a dilated local window around each query.
attentionsequence mixer
Native sparse attention: gated compression, selection, and sliding-window branches.
attentionfusedsequence mixer
Causal NSA attention over pre-compressed kv blocks; returns output and log-sum-exp.
attentionsequence mixer
Causal attention where keys are position-encoded by accumulated Householder reflections.
attentionpositionalsequence mixer
Polynomial composition of RMS-normalized powers of the input.
activationnormalization
Multimodal rotary embedding: temporal, height, and width tables interleaved per section.
positional
ReBased linear attention: a pure quadratic feature map on the scores.
linear attentionsequence mixer
Multi-scale retention: causal linear attention with a fixed per-head exponential decay.
linear attentionsequence mixer
Root-mean-square normalization.
normalization
RMS norm scaled by a swish or sigmoid gate.
activationfusednormalization
rms_norm_linear1,723 cases RMS norm fused with a linear projection.
feature mixerfusedlinearnormalization
rms_norm_linear_quant1,720 cases BitNet-style linear on an RMS-normed input: int8/ternary fake quant around the matmul.
fusedlinearnormalizationquantized
rms_norm_swish_linear1,723 cases RMS norm, swish gate, then a linear projection, in one fused op.
activationfeature mixerfusedlinearnormalization
Paired Llama-style rotary embedding of queries and keys.
positional
rotary_embedding898 cases Rotary position embedding, GPT-NeoX halves or GPT-J interleaved pairs.
positional
RWKV-4 wkv: a per-channel EMA of values weighted by exp(k), with a current-token bonus.
linear attentionsequence mixer
RWKV-6: linear attention with data-dependent per-key decay and a current-token bonus.
linear attentionsequence mixer
RWKV-7: a diagonal-plus-low-rank state transition with in-context erase and replace.
linear attentionsequence mixer
The RWKV-7 token-shift interpolation branches, fused.
feature mixerfused
rwkv7_channel_mixing1,291 cases RWKV-7 channel mixing: token-shift interpolation, squared-ReLU MLP, final-token state out.
activationfeature mixerfusedlinear
rwkv7_gate_output702 cases RWKV-7's receptance-key correction and output gate.
activationfused
rwkv7_k_update1,110 cases Interpolate RWKV-7 keys toward their in-context learning update.
feature mixerfused
selective_log_softmax1,343 cases Next-token log probabilities of the selected ids, computed in float32.
fusedreduction
Logistic sigmoid.
activation
Gated linear attention with a scalar per-head log forget gate.
linear attentionsequence mixer
Softmax over the last dimension.
activationreduction
Invert each block of a chunked unit-lower-triangular matrix.
linear
Sparse softmax: the Euclidean projection onto the probability simplex.
activationreduction
SwiGLU gating.
activationfeature mixerfused
SwiGLU gating fused with a linear projection.
activationfeature mixerfusedlinear
The LLaMA MLP: gate and up projections, SwiGLU, then the down projection.
activationfeature mixerfusedlinear
Titans linear memory: momentum gradient descent on a reconstruction loss, with forgetting.
linear attentionsequence mixer
RWKV token shift: the previous token's features minus the current ones.
sequence mixer
TTT-linear: the hidden state is a linear model trained by gradient descent at test time.
linear attentionsequence mixer
Total variation distance to a constant target distribution.
loss
Causal attention with per-channel decay on every logit, optional scalar gate and sink.
attentionsequence mixer
