Commit eef3aae

committed

feat: support custom mask for flash attention

1 parent b602c04 commit eef3aaeCopy full SHA for eef3aae

3 files changed

+274

-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

+274

-17

lines changed

Comments

(0)