A flexible and efficient implementation of Flash Attention 2.0 for JAX, supporting multiple backends (GPU/TPU/CPU) and platforms (Triton/Pallas/JAX).
-
Updated
Mar 4, 2025 - Python
#
pallas
Here are 15 public repositories matching this topic...
A FlashAttention backwards-over-backwards ⚡🔙🔙
-
Updated
Mar 28, 2025 - Jupyter Notebook
Calculate the hash of any input for ZK-Friendly hashes (MiMC & Poseidon) over a variety of Elliptic Curves.
-
Updated
Feb 7, 2023 - TypeScript
dLLM training implementation on pure jax/flax (w/o pytorch) for Google TPUs(v4/v5e/v6e). #TPUSprint #TRC
-
Updated
Apr 25, 2026 - Python
Benchmarking the JAX Pallas implementation of a custom RNN against alternatives
-
Updated
Feb 27, 2025 - Python
Packet-Switched Attention for stable 2-bit quantized MoE inference, with variance-aware routing and Protocol C benchmarks.
flax quantization pallas fine-tuning efficient-inference tpu jax xla mixture-of-experts memory-optimization sparse-attention opensource-ai gemma-2 llm-efficiency protocol-c
-
Updated
May 14, 2026 - Python
Flash Attention from first principles on TPU using JAX Pallas.
-
Updated
Apr 28, 2026 - Python
Repo to hold core components when building a Pallas Systems Website
-
Updated
Apr 1, 2025 - TypeScript
术 (Shu) — The first GPU-accelerated MSM for the Pallas curve. Part of the HanFei 韩非 series.
-
Updated
Mar 26, 2026 - Rust
Lean 4 formalization of the Pasta curves (Pallas and Vesta) for Zcash's Halo 2 — primality proofs and IsElliptic instances
-
Updated
May 4, 2026 - Lean
SuperNova (Pasta) proof generator & verifier with CI and frozen fixtures
-
Updated
Nov 24, 2025 - Rust
Improve this page
Add a description, image, and links to the pallas topic page so that developers can more easily learn about it.
Add this topic to your repo
To associate your repository with the pallas topic, visit your repo's landing page and select "manage topics."