Skip to content

FJFormer

Jax Flash Attn Gpu

FJFormer

Bits
Bits
- Bits
- Calibration
- Config
- Int Numerics
- No Numerics
- Numerics
- Q Dot General
- Q Flax
- Qk
- Stochastic Rounding
Checkpoint
Checkpoint
- Load
- Streamer
Func
Func
- Func
- Loss Func
Home
Monitor
Monitor
- Tracker
Optimizers
Optimizers
- Adafactor
- Adamw
- Lion
- Optimizer Utils
- Rmsprop
Pallas Operations
Pallas Operations
- Efficient Attention
  Efficient Attention
  - Efficient Attention
- Flash Attention
  Flash Attention
  - Gpu
    Gpu
    
    Jax Flash Attn Gpu Jax Flash Attn Gpu
    Table of contents
    
    jax_flash_attn_gpu
  - Tpu
    Tpu
    
    Jax Flash Attn Tpu
- Layer Norm
  Layer Norm
  - Gpu
    Gpu
    
    Layer Norm
- Ring Attention
  Ring Attention
  - Ring Attention
- Rms Norm
  Rms Norm
  - Gpu
    Gpu
    
    Rms Norm
- Softmax
  Softmax
  - Gpu
    Gpu
    
    Softmax
- Splash Attention
  Splash Attention
  - Tpu
    Tpu
    
    Splash Attention Kernel
    
    Splash Attention Mask
    
    Splash Attention Mask Info
Partition Utils
Partition Utils
- Mesh Utils
- T5x Partitioning
Utils
Xrapture
Xrapture

pallas_operations.flash_attention.gpu.jax_flash_attn_gpu

Module containing fused attention forward and backward pass.