atomic_ops: Fused Gated DeltaNet-2 Kernels for TPU v5e in JAX/Pallas
Year: 2026
Type: software
Abstract: A from-scratch port of the NVlabs Gated DeltaNet-2 Triton kernels to jax.experimental.pallas, targeting TPU v5e-8. Fused forward + backward custom_vjp kernels; on the full training shape (B=8, L=4096) achieves 27.2x (FP32) / 13.3x (BF16) speedup vs associative_scan baseline and 2.6x/3.4x vs pure-JAX WY reference
Author: Omirbay Akseleu
Institution: Nazarbayev Intellectual Schools
Subfield: Hardware and Architecture
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
OpenAlex ID: W7212303551
Cites:
Cited by:
Related to: