Frontier-class open models on a free Kaggle TPU v5e-8: GLM-5.3-Flash 320B MoE (~165 tok/s, our own JAX engine) and Qwen3.8-27B bf16 (~130 tok/s), 262k context, prefix caching. Works with Claude Code, Codex, opencode and pi.
A flexible and efficient implementation of Flash Attention 2.0 for JAX, supporting multiple backends (GPU/TPU/CPU) and platforms (Triton/Pallas/JAX).
A FlashAttention backwards-over-backwards ⚡🔙🔙
dLLM training implementation on pure jax/flax (w/o pytorch) for Google TPUs(v4/v5e/v6e). #TPUSprint #TRC
Calculate the hash of any input for ZK-Friendly hashes (MiMC & Poseidon) over a variety of Elliptic Curves.
Measurement harness for the sliding window attention premium in the vLLM TPU Ragged Paged Attention v3 kernel: per layer decode cost, block size control, throughput, and goodput for Gemma 4 31B on TPU v6e.
Google DeepMind's WeatherNext 2 at full 0.25° on consumer GPUs. Same model, same results: 6.4 GiB instead of 34 GiB, 7× faster.
PallasBench: Benchmarking Custom Kernel Generation in JAX Pallas. 150 tasks across 3 levels, fast_p metric, TPU+GPU. Inspired by KernelBench (Stanford, ICML 2025).
A GPU-accelerated pseudo-spectral solver for direct numerical simulation of the 3D incompressible Navier–Stokes equations, written in JAX.
Packet-Switched Attention for stable 2-bit quantized MoE inference, with variance-aware routing and Protocol C benchmarks.
GPU pricing of a cancellable swap, three ways: CUDA, ISO C++ stdpar, and JAX/XLA — Mark Joshi's kooderive modernized, rewritten twice, and cross-validated (LIBOR Market Model, Longstaff-Schwartz, AAD Greeks)
Hoata is the compiler explorer for TPU
Benchmarking the JAX Pallas implementation of a custom RNN against alternatives
A performant and compact LLM pretraining stack in pure JAX.
tpuasm 是 TPU 机器指令包的汇编器与反汇编器
SuperNova (Pasta) proof generator & verifier with CI and frozen fixtures
Repo to hold core components when building a Pallas Systems Website
术 (Shu) — The first GPU-accelerated MSM for the Pallas curve. Part of the HanFei 韩非 series.
To associate your repository with the pallas topic, visit your repo's landing page and select "manage topics."