Skip to content

Commit 70ddf77

Browse files
Update examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
1 parent eef32b3 commit 70ddf77

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ def main(
150150
scores_scale = T.alloc_fragment([block_M], accum_dtype)
151151
scores_sum = T.alloc_fragment([block_M], accum_dtype)
152152
logsum = T.alloc_fragment([block_M], accum_dtype)
153-
sinks = T.alloc_fragment([heads], dtype)
153+
sinks = T.alloc_fragment([block_M], dtype)
154154

155155
T.copy(Q[bz, by, bx * block_M:(bx + 1) * block_M, :], Q_shared)
156156
T.fill(acc_o, 0)

0 commit comments

Comments
 (0)