Sequence-mixer, kernel-sum and MoE performance¶

probjax's non-softmax kernels against the alternatives in the library and outside it, on GPU and CPU: SSD, the Mamba selective scan, residual linear attention, Gated DeltaNet, kernel matrix-vector products and KDE, log-domain kernel reductions (Sinkhorn, OT), and the MoE grouped matmul. GPU: RTX 3090; CPU: shared 24-thread host, whose timings carry 10–30% run-to-run noise. Attention is on its own page, setup and methodology on the overview. fwd+bwd is the forward plus the gradient with respect to all array inputs. Times are medians in milliseconds. OOM: out of device memory. fail†: XLA found no kernel configuration that fits during autotuning (NOT_FOUND), so the case did not run. unsup.: the implementation rejects the shape or backend. error: another failure (the JSON row has the message). —: not measured.

SSD (Mamba-2 state-space duality)¶

ssd(q, k, v, log_alpha) computes the Mamba-2 / SSD recurrence. implementation=None runs a chunked pure-JAX form (the Mamba-2 chunked algorithm with remat='full') on CPU and GPU; implementation='pallas' selects the Triton kernel, and ssd_linear_scan is the per-step reference. SSDCell and Mamba2Cell are built on ssd (models page).

Batch 4, 2 groups, 8 heads, d_k = d_v = 128, fp32 (kernel_ssd_gpu.json, kernel_ssd_cpu.json).

Device peak at sequence length 8192 (fresh process, MB):

Longest power-of-two sequence whose fwd+bwd ran on the 24 GB card (g2h8_d128):

Accuracy, relative (max|a − b| / max|b|) against an fp32 HIGHEST-precision run of the chunked algorithm: the chunked default's gradients differ by at most 1.3e-07 at sequence length 8192; the Pallas kernel runs its matmuls in TF32, with output error 5.2e-04 and gradient error up to 1.1e-02.

When to use which. Keep the default. On GPU the chunked form is at least as fast as the Triton kernel at the shapes above, keeps no residuals and runs any shape; it is the only fast path on CPU. ssd_linear_scan is a reference: its backward stores every state and runs out of memory at long sequences. ssd() forwards chunk_size, remat, precision and return_final_state to the chunked path.

Mamba selective scan¶

compute_mamba_scan is the selective scan of MambaCell. On GPU it runs Triton sequence-loop kernels that keep the state in registers and recompute instead of storing states in the backward pass; on CPU, and on GPU for unsupported shapes, a chunked XLA scan with a custom VJP. The associative scan and lax.scan are the generic JAX alternatives.

Batch 4, D 512, state 16, fp32; on CPU the associative scan was not run at 8192 (kernel_mamba_gpu.json, kernel_mamba_cpu.json).

Longest power-of-two sequence whose fwd+bwd ran on GPU (batch 4, D 1024, state 16; search limit 262,144):

GPU device peak, batch 4, D 512 (fresh process, MB):

The figure uses state 16; MambaCell defaults to state_dim = model_dim. GPU time of the Triton scan at larger states, median of three rounds (kernel_mamba_gpu_state.json), ms:

Accuracy. The fp32 implementations agree with lax.scan to 7.8e-05 relative error at sequence length 8192 (fwd+bwd, D 512).

When to use which. Keep the default dispatch. On GPU the Triton scan is the fastest measured implementation at every shape. On CPU the chunked scan is the fastest for fwd+bwd; for a forward-only pass lax.scan is within host noise of it and sometimes faster at D 512. The associative scan materializes [B, L, D, S] and lax.scan stores every state for the backward pass; both run out of memory well before the defaults.

Residual linear attention (RAttention)¶

residual_linear_attention, the linear-attention branch of RAttention, defaults to a chunk-parallel pure-JAX form (implementation='auto') on every backend. implementation='pallas' runs a Pallas forward with the chunked backward, implementation='scan' the token-by-token recurrence.

Batch 2, 4 heads, D 64, chunk 128, fp32; device peak at batch 4, 8 heads (kernel_rla_gpu.json, kernel_rla_cpu.json).

RAttention layer end to end (forward, and gradient of a scalar loss), batch 2, 8 heads, head_dim 64, window 128. The layer runs out of memory at 8192 tokens on the 24 GB GPU because its sliding-window branch is dense:

Accuracy, relative to the fp32 lax.scan: chunked 1.6e-06 (forward) and 4.9e-04 (fwd+bwd, max over outputs and gradients) at sequence length 8192, D 64; the Pallas forward 3.1e-04 (TF32).

When to use which. Keep implementation='auto': the chunked form is the fastest measured option on both backends and its backward does not store per-token state the way the scan backward does.

Gated DeltaNet¶

gated_delta_rule(implementation='auto') runs gated_delta_rule_pallas on GPU for supported chunk and head sizes, and the chunked pure-JAX form (gated_delta_rule_chunked, WY/UT representation with a hand-written VJP that keeps only chunk-boundary states) otherwise and on CPU. gated_delta_rule_recurrent is the token-by-token reference. ssd on the same shape is plotted as a cost reference; it has no delta-rule correction term. GatedDeltaNetCell is built on it.

bf16, batch 4, 8 heads, chunk 64, d_k = d_v = d, RTX 3090 (kernel_delta_gpu.json).

CPU, fp32, batch 2, 8 heads (kernel_delta_cpu.json; the recurrence was not run at 8192), ms:

Accuracy, relative to an fp32 HIGHEST-precision chunked reference: bf16 Pallas 4.5e-03 (fwd+bwd, d 128, sequence length 8192), fp32 Pallas 5.4e-04.

When to use which. Use the default. The recurrence stores every state in its backward pass and fails at long sequences. Compared with ssd on the same shape, Gated DeltaNet is cheaper at d 128 and more expensive at d 64 on this GPU; choose between them for the model, not the kernel.

RBF kernel matrix-vector products and KDE¶

rbf_kernel_mv computes K(q, k) @ v without materializing K; kde_density and rbf_kde_density build on it. The default is the Pallas kernel on GPU and a row-blocked XLA implementation on CPU (implementation='xla' on either).

fp32, D 8, 16 output columns, batch 1; GPU device peak in a fresh process (kernel_kmv_gpu.json, kernel_kmv_cpu.json).

GPU kde_density, ms:

Accuracy. Each JSON row carries rel_err, the norm-relative error against the row-blocked XLA implementation at HIGHEST matmul precision; the Pallas kernel is within 3.6e-04 (fwd+bwd, D 8, N = 16384, largest gradient error).

When to use which. Keep the defaults. At N = 1024 the dense product and the XLA path are faster than the Pallas kernel in fwd+bwd; from a few thousand points on the kernel is the fastest, and the dense product runs out of memory around N = 32k on the 24 GB GPU. The Pallas kernel ran fwd+bwd at N = 1,048,576 with a 320 MB device peak.

Log-domain kernel reductions, Sinkhorn and OT¶

kernel_logsumexp computes logsumexp_j(-c(x_i, y_j) / ε + g_j) without materializing the N × M cost: a Pallas kernel on GPU (cross term in 3×TF32), blocked XLA on CPU, dense for small problems. probjax.nn.ot builds on it (Sinkhorn with implicit or unrolled gradients, Sinkhorn divergence, entropic map, plan sampler, MMD, KDE log-density, ot_pairing); FlowMatcher / LinearFlow accept ot_coupling=None | 'exact' | 'sinkhorn'. The comparison is a dense jax.nn.logsumexp on the materialized cost (ott-jax is not installed).

fp32, N = M (kernel_ot_gpu.json, kernel_ot_cpu.json).

GPU device peak (fresh process):

Sinkhorn, 50 log-domain iterations, D 8, against a dense log-domain solver on the materialized cost; the last column is the absolute difference of the two transport costs:

Accuracy against an fp64 reference, N = M = 2048. dense_default_precision is the dense fp32 computation at JAX's default (TF32) matmul precision, which is why the kernel computes the cross term in 3×TF32:

OT-CFM (LinearFlow from N(0, I) to 8 Gaussians, 3000 steps, batch 256, GPU). Straightness is lower for straighter paths; the Sinkhorn divergence to held-out target samples after 1, 2, 4 and 32 Euler steps measures sample quality (the floor is the divergence between two target samples):

When to use which. On GPU the kernel matches or beats the dense forward from a few thousand points on and is the only option once the cost matrix no longer fits (dense fails from N = 32k–49k in the sweeps above). Its backward costs several times its forward, so for fwd+bwd at large D the dense form is faster while it fits (D 32, N = 16384: dense 12.1 ms vs kernel 15.0 ms, with 2,064 vs 0.188 MB of temporaries). For OT-CFM, Sinkhorn pairing costs little more per step than independent pairing and approaches the straightness of exact OT.

MoE grouped matmul¶

MoELayer(routing='dropless') sorts token-expert assignments by expert and runs one grouped matmul. On GPU the default is the Pallas grouped GEMM (dropless_impl='auto', dropless_dispatch='sorted'), which processes the grouped SwiGLU in row slabs so the train step holds no full-size intermediates; on CPU it is dropless_impl='chunked', an XLA loop over expert-sorted row chunks. routing='capacity' drops tokens beyond the expert capacity and is not exact; routing='dense' runs every expert on every token.

32768 tokens, d_model 512; 8 experts with d_hidden 1024, 64 experts with d_hidden 256; RTX 3090 (kernel_gmm_gpu.json).

Left: GPU device peak of the bf16 train step (8 experts top-2, fresh process). Right: CPU fp32 train step, 8 experts top-2 (kernel_gmm_cpu.json).

Largest bf16 train step that ran on the 24 GB GPU:

Accuracy, relative output error against an fp32 dense-routing reference: 1.1e-02 for fp32 Pallas on GPU (TF32 matmuls), 3.2e-02 in bf16, 3.6e-08 for the CPU chunked path.

When to use which. Keep the defaults. On GPU with 8 experts the sorted Pallas dispatch has the fastest exact train step and fits the most tokens, dropless_dispatch='fused_gsum' the fastest forward; with many small experts (64, top-4) fused_gsum also wins the train step. jax.lax.ragged_dot is far slower on this GPU. On CPU chunked is the fastest exact path.

How to reproduce¶

# Per-family runners (CUDA GPU); add JAX_PLATFORMS=cpu and --device cpu for the CPU files.
python -m docs.scripts.benchmark_kernel_ssd --device gpu --output docs/benchmarks/kernel_ssd_gpu.json
python -m docs.scripts.benchmark_kernel_mamba --device gpu --output docs/benchmarks/kernel_mamba_gpu.json
python -m docs.scripts.benchmark_kernel_rla --device gpu --suites kernel layer \
  --output docs/benchmarks/kernel_rla_gpu.json
python -m docs.scripts.benchmark_kernel_delta --device gpu --output docs/benchmarks/kernel_delta_gpu.json
python -m docs.scripts.benchmark_kernel_kmv --device gpu --output docs/benchmarks/kernel_kmv_gpu.json
python -m docs.scripts.benchmark_kernel_ot --device gpu --suite lse
python -m docs.scripts.benchmark_kernel_gmm --device gpu --output docs/benchmarks/kernel_gmm_gpu.json
JAX_PLATFORMS=cpu python -m docs.scripts.benchmark_kernel_ssd --device cpu \
  --output docs/benchmarks/kernel_ssd_cpu.json

# Rebuild the figures and this page from the JSON.
python -m docs.scripts.benchmark_figures

The Mamba runner's --peak and --max-len flags, and the peak, scale and cfm suites of the OT runner, produce the memory and largest-size rows.