Commit 44bb4e7

committed

feat: support custom mask for flash attention

1 parent b602c04 commit 44bb4e7Copy full SHA for 44bb4e7

3 files changed

+273

-17

lines changed

python/sgl_jax
- srt/layers/attention
  - flash_attn_kernel
    - flash_attention.py
  - flashattention_backend.py
- test
  - test_flashattention.py

3 files changed

+273

-17

lines changed

Comments

(0)