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: 561130379
  • Loading branch information
@wang12tao
wang12tao authored and jax authors committed Aug 29, 2023
1 parent 6072d59 commit 5a578cb
Show file tree
Hide file tree
Showing 2 changed files with 403 additions and 119 deletions.

0 comments on commit 5a578cb

Please sign in to comment.