Skip to content

Commit

Permalink
Add segment_ids support to pallas flash attention on GPU.
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 547592877
  • Loading branch information
@wang12tao
wang12tao authored and The jax_triton Authors committed Aug 25, 2023
1 parent 321a19e commit 5ca42f9
Showing 1 changed file with 2,019 additions and 0 deletions.
Loading

0 comments on commit 5ca42f9

Please sign in to comment.