Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[FA] fix an assertion failure due to refactoring in PR54 #69

Closed
wants to merge 3 commits into from
Closed

[FA] fix an assertion failure due to refactoring in PR54 #69

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c9e53ab
[FA] fix an assertion failure due to refactoring in PR54
manman-ren Nov 21, 2024
4921007
fix for _ws
manman-ren Nov 21, 2024
61b4299
add more configs so smaller headdim can pass
manman-ren Nov 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 15 additions & 12 deletions tritonbench/kernels/triton_fused_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -360,11 +360,11 @@ def _attn_fwd_inner_ws(
num_warps=w,
)
)
for BM in [128]
for BN in [128]
for BM in [64, 128]
for BN in [64, 128]
for sched in schedList
for enable_tma in [False]
for w in [8]
for w in [4, 8]
]
# no WS, with TMA and CompPipe
configsTma = [
Expand Down Expand Up @@ -393,11 +393,11 @@ def _attn_fwd_inner_ws(
num_warps=w,
)
)
for BM in [128]
for BN in [128]
for BM in [64, 128]
for BN in [64, 128]
for sched in schedList
for enable_tma in [True]
for w in [8]
for w in [4, 8]
]
# no TMA, with WS and CompPipe
configsWS = [
Expand Down Expand Up @@ -454,10 +454,10 @@ def _attn_fwd_inner_ws(
num_warps=w,
)
)
for BM in [128] # 64, 128]
for BN in [128] # 64, 128]
for s in [3] # , 4, 7]
for w in [8] # 4, 8]
for BM in [64, 128]
for BN in [64, 128]
for s in [3, 4, 7]
for w in [4, 8]
]
# TMA, WS, and CompPipe
configsTmaWS = [
Expand Down Expand Up @@ -548,7 +548,6 @@ def _attn_fwd_compute(
ENABLE_TMA: tl.constexpr,
LOOP_SCHEDULE: tl.constexpr,
):
tl.static_assert(BLOCK_N <= HEAD_DIM)
start_m = tl.program_id(0)
off_hz = tl.program_id(1)
off_z = off_hz // H
Expand Down Expand Up @@ -729,7 +728,6 @@ def _attn_fwd_compute_ws(
ENABLE_TMA: tl.constexpr,
LOOP_SCHEDULE: tl.constexpr,
):
tl.static_assert(BLOCK_N <= HEAD_DIM)
start_m = tl.program_id(0)
off_hz = tl.program_id(1)
off_z = off_hz // H
Expand Down Expand Up @@ -914,6 +912,7 @@ def _attn_fwd_ws(
LOOP_SCHEDULE: tl.constexpr,
ENABLE_WS: tl.constexpr,
):
tl.static_assert(BLOCK_N <= HEAD_DIM)
_attn_fwd_compute_ws(
Q,
K,
Expand Down Expand Up @@ -993,6 +992,7 @@ def _attn_fwd(
LOOP_SCHEDULE: tl.constexpr,
ENABLE_WS: tl.constexpr,
):
tl.static_assert(BLOCK_N <= HEAD_DIM)
_attn_fwd_compute(
Q,
K,
Expand Down Expand Up @@ -1072,6 +1072,7 @@ def _attn_fwd_opt( # Q, V, desc_k, desc_v, sm_scale, M, Out, #
LOOP_SCHEDULE: tl.constexpr,
ENABLE_WS: tl.constexpr,
):
tl.static_assert(BLOCK_N <= HEAD_DIM)
_attn_fwd_compute(
Q,
K,
Expand Down Expand Up @@ -1151,6 +1152,7 @@ def _attn_fwd_tma( # Q, V, desc_k, desc_v, sm_scale, M, Out, #
LOOP_SCHEDULE: tl.constexpr,
ENABLE_WS: tl.constexpr,
):
tl.static_assert(BLOCK_N <= HEAD_DIM)
_attn_fwd_compute(
Q,
K,
Expand Down Expand Up @@ -1230,6 +1232,7 @@ def _attn_fwd_tma_ws( # Q, V, desc_k, desc_v, sm_scale, M, Out, #
LOOP_SCHEDULE: tl.constexpr,
ENABLE_WS: tl.constexpr,
):
tl.static_assert(BLOCK_N <= HEAD_DIM)
_attn_fwd_compute_ws(
Q,
K,
Expand Down
Loading