Skip to content

GitLab

Unverified Commit 792e5d5b authored by Zhengju Tang's avatar Zhengju Tang Committed by GitHub
Browse files

[Feature] Add GQA backward kernel with varlen input (#1082) 

* [Feature] Add GQA backward kernel with varlen input

* [Lint]

* [BugFix] Freeze the memory order of all atomic_add operations

* [Lint]

* [Lint]

* [BugFix] Use release order to boost performance
parent bb8b3cd7
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 792 additions and 2 deletions
examples/flash_attention/example_gqa_bwd_tma_reduce_varlen.py 0 → 100644
View file @ 792e5d5b
This diff is collapsed.
examples/flash_attention/example_gqa_fwd_varlen.py
View file @ 792e5d5b
...@@ -8,8 +8,6 @@ from einops import rearrange, repeat ...@@ -8,8 +8,6 @@ from einops import rearrange, repeat
from tilelang.profiler import do_bench from tilelang.profiler import do_bench
from varlen_utils import generate_random_padding_mask, generate_qkv from varlen_utils import generate_random_padding_mask, generate_qkv
tilelang.disable_cache()
def attention_ref( def attention_ref(
q, q,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment