Skip to content

Commit 0704516

Browse files
bottlergemini-code-assist[bot]
authored andcommitted
test_attention compat with coming xformers change (vllm-project#20487)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com>
1 parent 0a1d53e commit 0704516

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

tests/kernels/attention/test_attention.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -450,7 +450,8 @@ def test_multi_query_kv_attention(
450450
start += seq_len
451451
# xformers.AttentionBias to Tensor for use in reference impl.
452452
alibi_bias = [
453-
b.materialize(b.shape, device=device).squeeze() for b in attn_bias
453+
b.materialize((1, num_query_heads, i, i), device=device).squeeze()
454+
for b, i in zip(attn_bias, seq_lens)
454455
]
455456
else:
456457
attn_bias = BlockDiagonalCausalMask.from_seqlens(seq_lens)

0 commit comments

Comments
 (0)