Op coverage

All ops are device-neutral and declared in the per-category headers under <brotensor/ops/> (ls include/brotensor/ops/ is the table of contents; <brotensor/ops.h> is an umbrella that includes them all).

Per-header surface

Header Surface
activation.h relu / tanh / sigmoid, silu, gelu (tanh-approx / exact / quick), GEGLU / GEGLU-exact / SwiGLU, snake (BigVGAN/DAC), elu (EnCodec), leaky_relu (HiFi-GAN)
attention.h single-head attention, MHA (optional biases), self/cross attention (train + flash), cross-attention with head-avg map + logit bias, attention token moments, self-attention with T5/ALiBi additive bias, Transformer-XL rel-pos bias producer (Conformer / FastConformer), SAM/ViTDet decomposed-rel-pos (incl. windowed), W8A16 bias-attention
flash_attention.h tiled flash attention (+ bare-core bwd), GQA prefill (causal or bidirectional), windowed (sliding-window causal), packed var-length (+ bwd), fused QKV+O projections (+ bwd), project-KV / Q-with-cached-KV, KV-cache append, causal flash-decode (GQA, + per-key-masked variant), W8A16 variants
linear.h linear (single / batched / fp16 / fused-act-epilogue), matmul (+ bwd), matmul_abt (batched strided A @ Bᵀ, 16-bit, fused bias+act epilogue), W8A16 batched linear
lora.h header-only LoRA adapter — low-rank delta on a frozen base linear, optional per-rank gate, forward + backward. Pure composition of public ops; runs on any backend
lstm.h single-layer LSTM over a sequence (PyTorch nn.LSTM layout): training forward with gate/cell caches + full-BPTT backward
norm.h LayerNorm (single + batched ±caches), RMSNorm, GroupNorm, BatchNorm (train/infer/bwd), per-head L2-norm (Gated DeltaNet), NCHW channel L2-normalize, pixel_norm (StyleGAN normalize_2nd_moment)
conv.h / conv1d.h conv2d / conv3d (+ W8A16), conv_transpose2d, deform_conv2d (torchvision modulated deformable v2, fwd), and the 1D family (conv1d wrappers, pad1d, conv_transpose1d, causal_conv1d + streaming update) — all with backward where applicable
rope.h RoPE forward/backward, rope_apply (explicit cos/sin tables, head-shared) + bwd, rope_apply_perhead (per-head cos/sin tables), M-RoPE (Qwen-VL three-axis)
delta_rule.h Gated Delta Rule linear attention — chunked prefill + streaming step
diffusion.h AdaLN modulate / broadcast_mul, fused ResBlock (+ W8A16, + bwd), DDIM / Euler / DPM++ 2M sampler steps, sinusoidal timestep embedding
stylegan.h StyleGAN3-R generator primitives: modulated_conv2d (style modulation + demod + conv, fwd+bwd), upfirdn2d (up→FIR→down resampler, fwd+bwd), bias_act (fused bias+act+gain+clamp, fwd+bwd), filtered_lrelu (alias-free nonlinearity, fwd+bwd)
spatial.h pad2d, slice2d, unfold2d (neighborhood im2col), window partition/reverse (SAM), spatial 2×2 patch merge (Qwen-VL block-major + pixel_unshuffle channel-major), patch_unpack (DiT unpatchify), pixel_shuffle_upsample_2x (DC-AE up-shortcut), NCHW↔sequence transpose
resize.h 2× nearest/bilinear up (+ bwd) — 2× avg down lives in pooling.h — arbitrary-scale interp2d (nearest/bilinear/bicubic, half-pixel + align-corners), convex (RAFT) upsample, 1D resample
pooling.h masked mean-pool, 2× avg downsample, adaptive avg pool2d, max pool2d (+ index bwd)
embedding.h embedding lookup (+ scatter bwd), gather_rows, scatter_rows (overwrite) / scatter_rows_add (accumulate)
concat.h concat/split rows, batched column-block concat, NCHW channel concat (+ bwd), copy_d2d + copy_d2d_strided (pitched 2D device copy)
reduction.h sum_rows / sum_cols, argmax_rows, top_k_rows, rows_count_above (per-row two-threshold counts → INT32, SAM AMG stability score)
loss.h softmax (+ bwd), row-batched softmax_rows, softmax-xent (+ segment / fused / fused-batched), MSE (vec / scalar / per-sample), BCE-with-logits fused-batched
optim.h sgd_step, adam_step, xavier_init
elementwise.h add / scale / clamp / mul-inplace, axpby / add_scalar / add_channel_bias in-place, dtype cast (FP32↔FP16/BF16), log / exp / round (+ bwd), sin / cos / rsqrt (+ bwd), threshold_u8 (binary threshold → INT8 byte mask, SAM AMG)
sampling.h sample_logits (temperature / top-k / top-p / greedy) + sample_logits_into (caller-owned buffers, CUDA-graph-safe), Philox RNG (randn / rand_uniform / rand_bernoulli / randn_truncated)
spectral.h complex ops, FFT/iFFT, rFFT/irFFT (+ bwd), STFT/iSTFT (+ bwd)
codec.h VQ encode (RVQ codeword search) + FSQ quantize (NanoCodec), straight-through bwd
quant.h W8A16 host quantizer + matmul, GGUF Q4_K / Q6_K / Q8_0 dequant + fused GEMV + batched matmul
image.h per-channel image normalize, uint8 HWC → FP32 NCHW

Backend coverage

The CPU backend implements essentially the entire FP32 surface — forward and backward, including the diffusion samplers, flash attention, the audio family, the vision primitives, LSTM, and the StyleGAN family — as the simple, correct, autovectorize-friendly reference. CPU is FP32-only by design: it leaves the FP16 / BF16 / INT8-W8A16 / GGUF-quant vtable slots null, and the dispatcher throws "brotensor: <op>: not implemented on CPU" if you call one.

The CUDA and Metal backends add the FP16 (and BF16) precision paths, batched-inference variants, the W8A16 and GGUF block-quant kernels, and a handful of GPU-only fused kernels. The audio op family is FP32 on all three backends with per-family CPU↔GPU parity tests.

Vulkan (the AMD GPU path) fills 266 of 270 slots — its per-slot table is vulkan-coverage.md.

Metal is at near-total parity with CUDA. Every Metal vtable slot is registered. CUDA leaves one null — filtered_lrelu_backward, where CUDA takes the composite; Metal has a fused backward (FP32, no dB) as well as the fused forward. Everything else in the table below is registered on both GPU backends.

Two "ops" are not vtable entries but device-agnostic compositions of public ops, so they run on any backend automatically: LoRA (ops/lora.h, header-only) and the filtered_lrelu composite (src/filtered_lrelu.cpp — the path on CPU, for the configs the fused CUDA/Metal kernels don't cover, and the CUDA backward).

GPU backends (CUDA / Metal)

FP32 fwd/bwd columns below mirror the CPU surface; the FP16 column is the GPU-only precision path (BF16 follows FP16 where the kernel notes it). INT8-W8A16 and GGUF-quant ops are GPU-only.

Op FP32 fwd FP32 bwd FP16 fwd Notes
matmul ✓ ✓ ✓ plain row-major A @ B (no bias); dtype-dispatched FP32 + FP16 (FP32 accumulation); backward returns dA/dB (caller zeros, op accumulates; FP16 uses FP32 scratch + fold). Metal FP32/FP16/BF16 (matmul, matmul bwd, batched linear fwd/bwd incl. FP16-weight/FP32-activation, matmul_abt, 1x1 conv): one FP32-accumulate simdgroup GEMM, 64×64 tiles (half tiles BK 32 for FP16×FP16, float tiles BK 16 otherwise), any layout/shape, M≤8 skinny GEMV-style kernel (src/metal/gemm_fp32.mm); Metal matmul bwd accumulates in the GEMM epilogue (one rounding, no scratch)
matmul_abt — — ✓ batched strided A @ Bᵀ: C[b](M,N) = A[b](M,K) @ B[b](N,K)ᵀ with caller-supplied per-slice element strides; 16-bit only (FP16 or BF16), FP32 accumulation, optional fused bias + activation epilogue; backs attention QK^T/PV and Linear. C is not auto-resized — the caller pre-sizes and dtypes it. Metal FP16 / BF16: the simdgroup GEMM in src/metal/gemm_fp32.mm (M≤8 or M·N≥1024, masked so any shape is valid), naive fallback below that. CUDA WMMA equivalent
matmul_int8w_fp16 — — ✓ W8A16 weight-only matmul; INT8 weights + per-row FP32 scales, FP16 acts, FP32 accum
linear_forward_batched_int8w_fp16 — — ✓ W8A16 batched linear in (B,in)→(B,out) layout; fused FP16 bias add; mirrors linear_forward_batched_fp16 shape contract; WMMA fast path for K%8==0 (FP16 tensor cores with INT8→FP16 dequant on shared-mem load), tiled fallback otherwise
linear ✓ ✓ ✓ dense; FP32 single + batched (fwd/bwd), FP16 batched-inference and batched-train backward (dtype-dispatched, FP32 scratch + fold)
relu / tanh / sigmoid ✓ ✓ ✓ elementwise; GPU dtype-dispatched FP32/FP16/BF16 on both the unary and batched fwd+bwd paths (relu/tanh have batched variants)
silu / gelu ✓ ✓ ✓ tanh-approx GELU; dtype-dispatched (FP16 bwd accumulates in FP32)
gelu_exact ✓ ✓ ✓ 0.5*x*(1+erf(x/√2)), exact PyTorch/diffusers default
quick_gelu ✓ ✓ ✓ x * sigmoid(1.702*x), OpenAI CLIP activation
geglu ✓ ✓ ✓ gated GELU (SD FFN); FP32+FP16 fwd/bwd, dtype-dispatched
geglu_exact ✓ ✓ ✓ gated exact-GELU FFN, matches diffusers GEGLU
swiglu ✓ ✓ ✓ gated SiLU FFN (Llama-style); dtype-dispatched FP32 + FP16
add / scale / mul_inplace ✓ n/a ✓ dtype-dispatched
axpby / add_scalar / add_channel_bias (in-place) ✓ n/a ✓ in-place y = a*x + b*y, scalar broadcast add, and per-channel NCHW bias add; dtype-dispatched
clamp ✓ n/a ✓ in-place min/max, dtype-dispatched (VAE epilogue)
sin / cos / rsqrt ✓ ✓ ✓ elementwise (Fourier features, demod/pixel-norm reciprocal-sqrt); GPU dtype-dispatched FP32/FP16/BF16 (FP32 math)
threshold_u8 ✓ n/a ✓ binary threshold → INT8 byte mask (SAM AMG mask binarization); FP32/FP16 input
build_slot_mask ✓ n/a — device-side validity mask construction
softmax ✓ ✓ — masked, numerically stable
softmax_rows ✓ n/a ✓ row-batched stable softmax — one launch covers all rows, replacing a per-row softmax loop (attention score matrices). May run in place (Y == X)
layernorm ✓ ✓ ✓ FP32 single + batched-infer; FP16 batched-infer + backward (dtype-dispatched, FP32 scratch + fold for dGamma/dBeta)
rms_norm ✓ ✓ ✓ y = x * gamma / sqrt(mean(x²) + eps); dtype-dispatched FP32 + FP16 (FP16 bwd uses FP32 scratch + fold for dGamma)
group_norm ✓ ✓ ✓ NCHW, per-group stats; dtype-dispatched fwd+bwd (FP16 bwd accumulates in FP32)
pixel_norm ✓ ✓ ✓ StyleGAN normalize_2nd_moment row normalize; GPU dtype-dispatched FP32/FP16/BF16
attention (single-head) ✓ ✓ —
mha (multi-head) ✓ ✓ —
self_attention ✓ ✓ ✓ FP32 = training (caches exposed via _train); FP16 = flash inference
cross_attention ✓ ✓ ✓ FP32 = training (caches exposed via _train, rectangular Wk/Wv); FP16 = flash inference
flash_attention — ✓ ✓ tiled online-softmax, Lk-unbounded, optional causal; FP16 backward via recompute returns dQ/dK/dV (no fwd-time caches). Bare-core bwd enables LoRA training when projections live outside the attention call.
flash_attention_gqa ✓ — ✓ grouped-query self-attention over pre-projected Q/K/V, causal or bidirectional. Q carries num_q_heads, K/V carry num_kv_heads (must divide; equal == plain MHA). causal=false is full bidirectional attention over all Lk keys — the encoder prefill a bidirectional-ified decoder needs (LLM2Vec); causal=true reproduces flash_attention_forward's mask with GQA. Tiled online softmax, FP32/FP16/BF16, FP32 accumulation; optional length-Lk key mask. Forward-only, all three backends
flash_attention_qkvo — — ✓ (fwd) / ✓ (bwd) fused Q/K/V/O projections + biases; rectangular Wk/Wv for cross-attn; optional causal; verified at SD1.5 U-Net head_dims (40/80/160) and CLIP head_dim 64. FP16 backward via recompute (no fwd-time caches); registered on CUDA and Metal (and the FP32 path on CPU). W8A16 variant (flash_attention_qkvo_int8w_fp16) routes all four projections through linear_forward_batched_int8w_fp16; attention core stays FP16
flash_attention_varlen — ✓ (bwd) ✓ packed variable-length MHA (Qwen-VL window attn); per-sequence boundaries from cu_seqlens INT32 prefix-sum buffers; optional per-sequence causal; recompute-based backward returns dQ/dK/dV
flash_attention_windowed ✓ — — sliding-window causal self-attention (streaming codecs / decode); window ≤ 0 is unbounded causal; FP32 (CPU + GPU)
flash_attention_project_kv — — ✓ pre-project ctx → K/V for cached cross-attention (SD timesteps reuse). W8A16 variant available
flash_attention_q_with_kv_cached — — ✓ forward against pre-projected K/V; bitwise-equivalent to flash_attention_qkvo's cached path. W8A16 variant available
flash_attention_decode — — ✓ causal-aware decode against a partially-filled K/V cache; supports L_q ≥ 1 (token-by-token or chunked). Masked variant flash_attention_decode_masked (per-key FP32 validity mask, L_q == 1) on all three backends
kv_cache_append — — ✓ append L_new projected K/V rows into a pre-allocated L_max cache at cur_len
rope ✓ ✓ ✓ rotary position embedding; pair-wise rotation per head_dim chunk, seq_offset for KV-cache decode
resblock — ✓ (bwd) ✓ fused diffusion ResBlock (GN→SiLU→conv ×2 + skip); FP16 backward via composition of public ops (recomputes h1/h2/h3; no fwd-time caches)
conv2d ✓ ✓ ✓ NCHW, stride/pad/dil; groups ≥ 1 (depthwise supported); backward (dX, dW, dB) dtype-dispatched (FP32+FP16; FP16 dW/dB use FP32 scratch + fold)
conv2d_int8w_fp16 — — ✓ W8A16 weight-only conv2d; INT8 OIHW filter + per-output-channel FP32 scales, FP16 acts; CUDA WMMA fast path for 3x3 s1, 1x1 s1, 3x3 s2 (groups=1, dil=1) — naive fallback otherwise (incl. low-CTA long-K shapes)
deform_conv2d ✓ — ✓ torchvision deform_conv2d v2 (modulated deformable, bilinear per-tap offsets), forward-only; FP32/FP16, FP32 accumulation
modulated_conv2d ✓ ✓ ✓ StyleGAN synthesis core: per-sample style modulation + optional demodulation + conv2d; dW optional (skippable for inversion); GPU dtype-dispatched FP32/FP16/BF16, FP32 reductions
upfirdn2d ✓ ✓ ✓ upsample → pad/crop → 2D FIR → downsample (StyleGAN3, incl. non-separable config-R radial filters); backward is upfirdn2d with up/down swapped; GPU FP32/FP16/BF16, FP32 filter math
bias_act ✓ ✓ ✓ fused per-channel bias + activation (linear/lrelu) + gain + clamp; GPU FP32/FP16/BF16 (FP32 math, FP32-scratch dB)
filtered_lrelu ✓ ✓ ✓ alias-free nonlinearity (bias → upsample → lrelu → downsample). Forward: fused CUDA / Metal kernel where the config is covered and no up_buf/act_buf cache is requested, otherwise the device-agnostic composite over bias_act/upfirdn2d (always the path on CPU). Backward: fused on Metal (FP32, dB not requested; recomputes the pre-activation, so it needs no cache), the composite elsewhere
lstm ✓ ✓ — single-layer training LSTM + full BPTT (PyTorch nn.LSTM weight layout, gate order i|f|g|o); FP32 on all three backends
upsample_nearest_2x ✓ ✓ ✓ backward dtype-dispatched (FP32+FP16)
upsample_bilinear_2x ✓ ✓ ✓ align_corners=False; backward dtype-dispatched (FP32+FP16; FP16 uses FP32 scratch + fold)
downsample_avg_2x ✓ ✓ ✓ stride 2, kernel 2; backward dtype-dispatched (FP32+FP16)
nchw ↔ sequence transpose ✓ n/a ✓ gather/scatter between NCHW and (L,D) layouts
embedding lookup ✓ ✓ ✓ FP32/FP16 table dispatch; backward dtype-dispatched (FP16 uses FP32 scratch + fold for atomic-add safety)
concat_rows / split_rows ✓ ✓ ✓ flat byte-aware concat (FP16 supported)
concat_batched_rows ✓ n/a ✓ per-row column-block concat via 2D memcpy
concat_nchw_channels ✓ ✓ ✓ channel-axis concat for U-Net skip merges (N≥1); backward is per-part scatter (overwrites parts)
masked_mean_pool ✓ ✓ ✓ row-wise mean over valid mask; GPU dtype-dispatched FP32/FP16/BF16 (fwd + bwd)
sum_rows / sum_cols ✓ n/a ✓ reductions along rows/cols; dtype-dispatched FP32 + FP16
argmax_rows ✓ n/a ✓ per-row argmax; FP32/FP16 input, FP32 indices
rows_count_above ✓ n/a ✓ per-row above-threshold counts at two thresholds in one pass → (R,2) INT32 (SAM AMG stability score); FP32/FP16 input
ddim_step — n/a ✓ fused DDIM sampler step over FP16 latents; FP32 internal math
euler_step — n/a ✓ fused Euler-discrete step (ε-prediction, σ convention; matches diffusers EulerDiscreteScheduler)
dpmpp_2m_step — n/a ✓ fused DPM-Solver++ 2M multistep update; caller supplies linear-combo coefficients and x0 cache. First step falls back to euler_step
timestep_embedding ✓ n/a — sinusoidal embedding (FP32) for diffusion timesteps and SDXL added-cond micro-conditioning; diffusers default (flip_sin_to_cos=True)
copy_d2d ✓ n/a ✓ flat-buffer device-to-device chunk copy. Strided variant copy_d2d_strided (pitched 2D copy, e.g. NCHW W-axis pad/unpad) on all three backends
build_causal_mask_row n/a n/a ✓ length-L FP32 mask, CLIP text
sgd / adam ✓ n/a — optimizer steps
mse / softmax-xent / bce ✓ ✓ — per-sample + batched (BCE-with-logits fused-batched added)
conv3d ✓ — ✓ NCTHW (forward only), grouped, FP32/FP16/BF16; W8A16 conv3d_int8w_fp16 for Qwen-VL patch-embed (GPU-only)
conv_transpose1d / 2d ✓ ✓ — learned upsample (vocoders, SAM mask decoder, DPT heads); FP32, all three backends
batch_norm ✓ ✓ ✓ (inference) NCHW train/infer/bwd, running stats (pretrained ResNet/DETR backbones); training fwd + bwd FP32-only; inference fwd dtype-dispatched FP32/FP16/BF16 (FP32 math per element)
l2_norm / l2_normalize_nchw ✓ ✓ (l2_norm) ✓ per-head q/k L2 (Gated DeltaNet) + channel-axis NCHW normalize (DSINE normals); FP32/FP16/BF16
gated_delta_rule ✓ n/a ✓ chunked prefill + streaming step (linear-attention text decoders); FP32/FP16, FP32 accumulators
rope_apply / rope_apply_perhead / rope_apply_mrope ✓ ✓ (apply) ✓ explicit cos/sin tables (2D axial RoPE), head-shared or per-head; + Qwen-VL three-axis M-RoPE; FP32/FP16/BF16 (perhead/mrope are inference-only, no bwd)
self_attention_bias ✓ n/a ✓ additive pre-softmax bias (T5 rel-pos / ALiBi); FP32/FP16/BF16; W8A16 variant GPU-only
decomposed_rel_pos (± windowed) ✓ n/a ✓ SAM/ViTDet data-dependent 2D rel-pos attention; FP32/FP16/BF16
modulate / broadcast_mul ✓ n/a ✓ AdaLN affine + per-channel gate (DiT/SD3/Flux); FP32/FP16/BF16
pad2d / slice2d / unfold2d ✓ ✓ (pad/slice) ✓ image pad (zero/reflect/replicate), crop, neighborhood im2col
window_partition / spatial_merge ✓ n/a ✓ SAM window tiling (+ reverse) and 2×2 patch merge (Qwen-VL block-major / pixel_unshuffle channel-major, e.g. Flux.2 VAE)
patch_unpack ✓ n/a ✓ DiT unpatchify — (hpwp, PP*C_total) tokens → NCHW, block-major or channel-major column composition; channels past C_keep dropped. Inference-only; FP32/FP16/BF16
pixel_shuffle_upsample_2x ✓ n/a ✓ fused repeat_interleave + pixel_shuffle(2) (DC-AE up-shortcut) with no intermediate 4·C_out tensor. C_in == 4·C_out is a plain pixel-shuffle; C_in == C_out reduces to 2× nearest. Pure gather, inference-only; FP32/FP16/BF16
interp2d / convex_upsample ✓ ✓ (interp) ✓ arbitrary-scale resize (nearest/bilinear/bicubic, half-pixel + align-corners) + RAFT convex upsample
max_pool2d / adaptive_avg_pool2d ✓ ✓ — indexed max-pool bwd + PyTorch adaptive avg pool; FP32, all three backends
gather_rows / scatter_rows / scatter_rows_add ✓ ✓ — index-driven row gather + scatter (overwrite / accumulate) (SAM prompt encoder, DETR queries); FP32, all three backends
top_k_rows ✓ n/a — per-row top-k values + indices; FP32, all three backends
randn / rand_uniform / rand_bernoulli / randn_truncated ✓ n/a — Philox 4×32-10 RNG, PyTorch/JAX-compatible, FP32, all three backends
GGUF Q4_K / Q6_K / Q8_0 — — ✓ block-quant dequant + fused GEMV + batched matmul (W4/6/8-A16); registered on CUDA and Metal. Batched prefill (large B) uses a tensor-core GEMM — CUDA fused-dequant WMMA, Metal dequant-once + simdgroup matmul; small B is per-token GEMV

Audio op family

An FP32 op family for TTS / STT / neural-codec inference, consumed by the brosoundml sibling. Like the rest of the FP32 surface these are implemented on all three backends — CPU, CUDA, and Metal — FP32 throughout, with per-family CPU↔GPU parity tests.

Group Ops
Spectral core fft / ifft, rfft / irfft (+ backward), complex_mul / complex_abs / complex_angle / complex_from_polar — interleaved-complex FP32, mixed-radix + Bluestein
STFT stft / istft (+ backward) — windowed, COLA-normalised overlap-add
1D convolution conv1d (+ 3 backward halves), pad1d, causal_conv1d, conv_transpose1d (vocoder upsampling), causal_conv1d_update (streaming state)
Vocoder / codec activations snake / snakebeta (BigVGAN/DAC), elu (EnCodec), leaky_relu (HiFi-GAN)
Codec quantization vq_encode (RVQ codeword search), fsq_quantize (NanoCodec FSQ) — straight-through backward
Resampling resample1d — arbitrary-scale length resample, nearest / linear (+ backward)
Elementwise log / exp / round (+ backward) — log-mel domain maps
Sampling sample_logits — temperature / top-k / top-p autoregressive sampler, counter-based Philox RNG. sample_logits_into is the same sampler over caller-owned scratch/counter/index buffers, so a decode step allocates nothing and stays CUDA-graph-capturable

The conv1d family is header-only wrappers over the conv2d ops (a 1D conv is a 2D conv with H=kH=1); pad1d, conv_transpose1d, and causal_conv1d_update are genuine per-backend kernels. Two exceptions to the FP32-everywhere rule: conv1d_int8w_fp16 is a GPU-only W8A16 wrapper like its conv2d_int8w_fp16 parent, and the CUDA leaky_relu forward is dtype-dispatched FP32/FP16/BF16 (FP32 math per element).

Tests

ctest --test-dir build -C Release

Tests live under tests/, enabled by BROTENSOR_TESTS=ON (default ON when built standalone).