Fix bug in varlen test

ROCm · May 20, 2024 · 4533a48 · 4533a48
1 parent 4c06dc8
commit 4533a48
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/python/perf-kernels/flash-attention.py b/python/perf-kernels/flash-attention.py
@@ -833,6 +833,7 @@ def _attn_bwd(Q, K, V, sm_scale, alibi_slopes, DO, DQ, DK, DV, M, D,
 def get_shape_from_layout(q, k, metadata):
     if metadata.layout == 'thd':
         nheads_q, nheads_k = q.shape[1], k.shape[1]
+        head_size = q.shape[-1]
         batch = metadata.num_contexts
     elif metadata.layout == 'bhsd':
         batch, nheads_q, _, head_size = q.shape